Merge pull request #229 from avik-pal/ap/concrete_jac

Fix Jacobian Construction for concrete_jac

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "2.0.1"
+version = "2.0.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -35,6 +35,7 @@ FiniteDiff = "2"
 ForwardDiff = "0.10.3"
 LineSearches = "7"
 LinearSolve = "2"
+NonlinearProblemLibrary = "0.1"
 PrecompileTools = "1"
 RecursiveArrayTools = "2"
 Reexport = "0.2, 1"
@@ -52,6 +53,7 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -62,4 +64,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools"]
+test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary"]

src/NonlinearSolve.jl

@@ -4,7 +4,7 @@ if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@max_m
     @eval Base.Experimental.@max_methods 1
 end
 
-using DiffEqBase, LinearAlgebra, LinearSolve, SparseDiffTools
+using DiffEqBase, LinearAlgebra, LinearSolve, SparseArrays, SparseDiffTools
 import ForwardDiff
 
 import ADTypes: AbstractFiniteDifferencesMode

src/jacobian.jl

@@ -1,4 +1,4 @@
-@concrete struct JacobianWrapper{iip}
+@concrete struct JacobianWrapper{iip} <: Function
     f
     p
 end
@@ -73,9 +73,9 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii
         jac_cache = nothing
     end
 
-    J = if !linsolve_needs_jac
+    J = if !(linsolve_needs_jac || alg_wants_jac)
         # We don't need to construct the Jacobian
-        JacVec(uf, u; autodiff = alg.ad)
+        JacVec(uf, u; autodiff = __get_nonsparse_ad(alg.ad))
     else
         if has_analytic_jac
             f.jac_prototype === nothing ? undefmatrix(u) : f.jac_prototype
@@ -98,6 +98,11 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii
     return uf, linsolve, J, fu, jac_cache, du
 end
 
+__get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff()
+__get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff()
+__get_nonsparse_ad(::AutoSparseZygote) = AutoZygote()
+__get_nonsparse_ad(ad) = ad
+
 ## Special Handling for Scalars
 function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u::Number, p,
     ::Val{false}; kwargs...)

src/levenberg.jl

@@ -216,7 +216,7 @@ function perform_step!(cache::LevenbergMarquardtCache{true})
     # The following lines do: cache.a = -J \ cache.fu_tmp
     mul!(cache.du, J, v)
     @. cache.fu_tmp = (2 / h) * ((cache.fu_tmp - fu1) / h - cache.du)
-    linres = dolinsolve(alg.precs, linsolve; A = J, b = _vec(cache.fu_tmp),
+    linres = dolinsolve(alg.precs, linsolve; A = cache.mat_tmp, b = _vec(cache.fu_tmp),
         linu = _vec(cache.du), p = p, reltol = cache.abstol)
     cache.linsolve = linres.cache
     @. cache.a = -cache.du
@@ -225,7 +225,7 @@ function perform_step!(cache::LevenbergMarquardtCache{true})
 
     # Require acceptable steps to satisfy the following condition.
     norm_v = norm(v)
-    if (2 * norm(cache.a) / norm_v) < α_geodesic
+    if 2 * norm(cache.a) ≤ α_geodesic * norm_v
         @. cache.δ = v + cache.a / 2
         @unpack δ, loss_old, norm_v_old, v_old, b_uphill = cache
         f(cache.fu_tmp, u .+ δ, p)
@@ -274,18 +274,19 @@ function perform_step!(cache::LevenbergMarquardtCache{false})
     end
     @unpack u, p, λ, JᵀJ, DᵀD, J = cache
 
+    cache.mat_tmp = JᵀJ + λ * DᵀD
     # Usual Levenberg-Marquardt step ("velocity").
-    cache.v = -(JᵀJ + λ * DᵀD) \ (J' * fu1)
+    cache.v = -cache.mat_tmp \ (J' * fu1)
 
     @unpack v, h, α_geodesic = cache
     # Geodesic acceleration (step_size = v + a / 2).
-    cache.a = -J \ ((2 / h) .* ((f(u .+ h .* v, p) .- fu1) ./ h .- J * v))
+    cache.a = -cache.mat_tmp \ ((2 / h) .* ((f(u .+ h .* v, p) .- fu1) ./ h .- J * v))
     cache.stats.nsolve += 1
     cache.stats.nfactors += 1
 
     # Require acceptable steps to satisfy the following condition.
     norm_v = norm(v)
-    if (2 * norm(cache.a) / norm_v) < α_geodesic
+    if 2 * norm(cache.a) ≤ α_geodesic * norm_v
         cache.δ = v .+ cache.a ./ 2
         @unpack δ, loss_old, norm_v_old, v_old, b_uphill = cache
         fu_new = f(u .+ δ, p)

test/23_test_problems.jl

@@ -0,0 +1,62 @@
+using NonlinearSolve, LinearAlgebra, NonlinearProblemLibrary, Test
+
+problems = NonlinearProblemLibrary.problems
+dicts = NonlinearProblemLibrary.dicts
+
+function test_on_library(problems, dicts, alg_ops, broken_tests, ϵ = 1e-5)
+    for (idx, (problem, dict)) in enumerate(zip(problems, dicts))
+        x = dict["start"]
+        res = similar(x)
+        nlprob = NonlinearProblem(problem, x)
+        @testset "$(dict["title"])" begin
+            for alg in alg_ops
+                sol = solve(nlprob, alg, abstol = 1e-15, reltol = 1e-15)
+                problem(res, sol.u, nothing)
+                broken = idx in broken_tests[alg] ? true : false
+                @test norm(res)≤ϵ broken=broken
+            end
+        end
+    end
+end
+
+# NewtonRaphson
+@testset "NewtonRaphson test problem library" begin
+    alg_ops = (NewtonRaphson(),)
+
+    # dictionary with indices of test problems where method does not converge to small residual
+    broken_tests = Dict(alg => Int[] for alg in alg_ops)
+    broken_tests[alg_ops[1]] = [1, 6]
+
+    test_on_library(problems, dicts, alg_ops, broken_tests)
+end
+
+@testset "TrustRegion test problem library" begin
+    alg_ops = (TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Simple),
+        TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Fan),
+        TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Hei),
+        TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Yuan),
+        TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Bastin),
+        TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.NLsolve))
+
+    # dictionary with indices of test problems where method does not converge to small residual
+    broken_tests = Dict(alg => Int[] for alg in alg_ops)
+    broken_tests[alg_ops[1]] = [6, 11, 21]
+    broken_tests[alg_ops[2]] = [6, 11, 21]
+    broken_tests[alg_ops[3]] = [1, 6, 11, 12, 15, 16, 21]
+    broken_tests[alg_ops[4]] = [1, 6, 8, 11, 15, 16, 21, 22]
+    broken_tests[alg_ops[5]] = [6, 21]
+    broken_tests[alg_ops[6]] = [6, 21]
+
+    test_on_library(problems, dicts, alg_ops, broken_tests)
+end
+
+@testset "TrustRegion test problem library" begin
+    alg_ops = (LevenbergMarquardt(), LevenbergMarquardt(; α_geodesic = 0.5))
+
+    # dictionary with indices of test problems where method does not converge to small residual
+    broken_tests = Dict(alg => Int[] for alg in alg_ops)
+    broken_tests[alg_ops[1]] = [3, 6, 11, 17, 21]
+    broken_tests[alg_ops[2]] = [3, 6, 11, 21]
+
+    test_on_library(problems, dicts, alg_ops, broken_tests)
+end

test/runtests.jl

@@ -15,6 +15,8 @@ end
     if GROUP == "All" || GROUP == "Core"
         @time @safetestset "Basic Tests + Some AD" include("basictests.jl")
         @time @safetestset "Sparsity Tests" include("sparse.jl")
+
+        @time @safetestset "23 Test Problems" include("23_test_problems.jl")
     end
 
     if GROUP == "GPU"