Distributed Testing using ReTestItems

Author: avik-pal

.github/workflows/CI.yml

@@ -44,7 +44,9 @@ jobs:
       - uses: julia-actions/julia-runtest@v1
         env:
           GROUP: ${{ matrix.group }}
-          JULIA_NUM_THREADS: 8
+          JULIA_NUM_THREADS: 11
+          RETESTITEMS_NWORKERS: 4
+          RETESTITEMS_NWORKER_THREADS: 2
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: src,ext

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.5.3"
+version = "3.5.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -84,6 +84,7 @@ PrecompileTools = "1.2"
 Preferences = "1.4"
 Printf = "1.10"
 Random = "1.91"
+ReTestItems = "1"
 RecursiveArrayTools = "3.4"
 Reexport = "1.2"
 SIAMFANLEquations = "1.0.1"
@@ -101,7 +102,6 @@ Symbolics = "5.13"
 Test = "1.10"
 TimerOutputs = "0.5.23"
 Zygote = "0.6.67"
-XUnit = "1.1"
 julia = "1.10"
 
 [extras]
@@ -123,6 +123,8 @@ NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
@@ -132,8 +134,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-XUnit = "3e3c03f2-1a94-11e9-2981-050a4ca824ab"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials", "XUnit"]
+test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration", "SIAMFANLEquations", "Sundials", "ReTestItems", "Reexport"]

src/algorithms/levenberg_marquardt.jl

@@ -69,10 +69,11 @@ end
     J_diag_cache
     J_damped
     damping_f
+    initial_damping
 end
 
 function reinit_cache!(cache::LevenbergMarquardtDampingCache, args...; kwargs...)
-    cache.λ = cache.damping_f.initial_damping
+    cache.λ = cache.initial_damping
     cache.λ_factor = cache.damping_f.increase_factor
     if !(cache.DᵀD isa Number)
         if can_setindex(cache.DᵀD.diag)
@@ -110,7 +111,7 @@ function __internal_init(prob::AbstractNonlinearProblem,
     J_damped = T(initial_damping) .* DᵀD
     return LevenbergMarquardtDampingCache(T(f.increase_factor), T(f.decrease_factor),
         T(f.min_damping), T(f.increase_factor), T(initial_damping), DᵀD, J_diag_cache,
-        J_damped, f)
+        J_damped, f, T(initial_damping))
 end
 
 (damping::LevenbergMarquardtDampingCache)(::Nothing) = damping.J_damped

test/core/23_test_problems.jl renamed to test/core/23_test_problems_tests.jl

@@ -1,4 +1,5 @@
-using NonlinearSolve, LinearAlgebra, LinearSolve, NonlinearProblemLibrary, XUnit
+@testsetup module RobustnessTesting
+using NonlinearSolve, LinearAlgebra, LinearSolve, NonlinearProblemLibrary, Test
 
 problems = NonlinearProblemLibrary.problems
 dicts = NonlinearProblemLibrary.dicts
@@ -36,7 +37,10 @@ function test_on_library(problems, dicts, alg_ops, broken_tests, ϵ = 1e-4;
     end
 end
 
-@testcase "NewtonRaphson 23 Test Problems" begin
+export test_on_library, problems, dicts
+end
+
+@testitem "NewtonRaphson" setup=[RobustnessTesting] begin
     alg_ops = (NewtonRaphson(),)
 
     broken_tests = Dict(alg => Int[] for alg in alg_ops)
@@ -45,7 +49,7 @@ end
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
 
-@testcase "TrustRegion 23 Test Problems" begin
+@testitem "TrustRegion" setup=[RobustnessTesting] begin
     alg_ops = (TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Simple),
         TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Fan),
         TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Hei),
@@ -64,7 +68,9 @@ end
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
 
-@testcase "LevenbergMarquardt 23 Test Problems" begin
+@testitem "LevenbergMarquardt" setup=[RobustnessTesting] begin
+    using LinearSolve
+
     alg_ops = (LevenbergMarquardt(),
         LevenbergMarquardt(; α_geodesic = 0.1),
         LevenbergMarquardt(; linsolve = CholeskyFactorization()))
@@ -77,7 +83,7 @@ end
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
 
-@testcase "DFSane 23 Test Problems" begin
+@testitem "DFSane" setup=[RobustnessTesting] begin
     alg_ops = (DFSane(),)
 
     broken_tests = Dict(alg => Int[] for alg in alg_ops)
@@ -86,7 +92,7 @@ end
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
 
-@testcase "Broyden 23 Test Problems" begin
+@testitem "Broyden" setup=[RobustnessTesting] begin
     alg_ops = (Broyden(),
         Broyden(; init_jacobian = Val(:true_jacobian)),
         Broyden(; update_rule = Val(:bad_broyden)),
@@ -101,7 +107,7 @@ end
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
 
-@testcase "Klement 23 Test Problems" begin
+@testitem "Klement" setup=[RobustnessTesting] begin
     alg_ops = (Klement(), Klement(; init_jacobian = Val(:true_jacobian_diagonal)))
 
     broken_tests = Dict(alg => Int[] for alg in alg_ops)
@@ -111,7 +117,7 @@ end
     test_on_library(problems, dicts, alg_ops, broken_tests)
 end
 
-@testcase "PseudoTransient 23 Test Problems" begin
+@testitem "PseudoTransient" setup=[RobustnessTesting] begin
     # PT relies on the root being a stable equilibrium for convergence, so it won't work on
     # most problems
     alg_ops = (PseudoTransient(),)

test/core/forward_ad.jl renamed to test/core/forward_ad_tests.jl

@@ -1,5 +1,6 @@
-using ForwardDiff,
-    NonlinearSolve, MINPACK, NLsolve, StaticArrays, Sundials, XUnit, LinearAlgebra
+@testsetup module ForwardADTesting
+using Reexport, NonlinearSolve
+@reexport using ForwardDiff, MINPACK, NLsolve, StaticArrays, Sundials, LinearAlgebra
 
 test_f!(du, u, p) = (@. du = u^2 - p)
 test_f(u, p) = (@. u^2 - p)
@@ -58,50 +59,55 @@ __compatible(::NLsolveJL, ::Val{:oop_cache}) = false
 __compatible(::KINSOL, ::Val{:iip_cache}) = false
 __compatible(::KINSOL, ::Val{:oop_cache}) = false
 
-@testcase "ForwardDiff.jl Integration: $(alg)" for alg in (NewtonRaphson(), TrustRegion(),
-    LevenbergMarquardt(), PseudoTransient(; alpha_initial = 10.0), Broyden(), Klement(),
-    DFSane(), nothing, NLsolveJL(), CMINPACK(),
-    KINSOL(; globalization_strategy = :LineSearch))
-    us = (2.0, @SVector[1.0, 1.0], [1.0, 1.0], ones(2, 2), @SArray ones(2, 2))
+export test_f!, test_f, jacobian_f, solve_with, __compatible
+end
+
+@testitem "ForwardDiff.jl Integration" setup=[ForwardADTesting] begin
+    for alg in (NewtonRaphson(), TrustRegion(),
+        LevenbergMarquardt(), PseudoTransient(; alpha_initial = 10.0), Broyden(), Klement(),
+        DFSane(), nothing, NLsolveJL(), CMINPACK(),
+        KINSOL(; globalization_strategy = :LineSearch))
+        us = (2.0, @SVector[1.0, 1.0], [1.0, 1.0], ones(2, 2), @SArray ones(2, 2))
 
-    @testset "Scalar AD" begin
-        for p in 1.0:0.1:100.0, u0 in us, mode in (:iip, :oop, :iip_cache, :oop_cache)
-            __compatible(u0, alg) || continue
-            __compatible(u0, Val(mode)) || continue
-            __compatible(alg, Val(mode)) || continue
+        @testset "Scalar AD" begin
+            for p in 1.0:0.1:100.0, u0 in us, mode in (:iip, :oop, :iip_cache, :oop_cache)
+                __compatible(u0, alg) || continue
+                __compatible(u0, Val(mode)) || continue
+                __compatible(alg, Val(mode)) || continue
 
-            sol = solve(NonlinearProblem(test_f, u0, p), alg)
-            if SciMLBase.successful_retcode(sol)
-                gs = abs.(ForwardDiff.derivative(solve_with(Val{mode}(), u0, alg), p))
-                gs_true = abs.(jacobian_f(u0, p))
-                if !(isapprox(gs, gs_true, atol = 1e-5))
-                    @show sol.retcode, sol.u
-                    @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_gradient=gs true_gradient=gs_true
-                else
-                    @test abs.(gs)≈abs.(gs_true) atol=1e-5
+                sol = solve(NonlinearProblem(test_f, u0, p), alg)
+                if SciMLBase.successful_retcode(sol)
+                    gs = abs.(ForwardDiff.derivative(solve_with(Val{mode}(), u0, alg), p))
+                    gs_true = abs.(jacobian_f(u0, p))
+                    if !(isapprox(gs, gs_true, atol = 1e-5))
+                        @show sol.retcode, sol.u
+                        @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_gradient=gs true_gradient=gs_true
+                    else
+                        @test abs.(gs)≈abs.(gs_true) atol=1e-5
+                    end
                 end
             end
         end
-    end
 
-    @testset "Jacobian" begin
-        for u0 in us, p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0]),
-            mode in (:iip, :oop, :iip_cache, :oop_cache)
+        @testset "Jacobian" begin
+            for u0 in us, p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0]),
+                mode in (:iip, :oop, :iip_cache, :oop_cache)
 
-            __compatible(u0, p) || continue
-            __compatible(u0, alg) || continue
-            __compatible(u0, Val(mode)) || continue
-            __compatible(alg, Val(mode)) || continue
+                __compatible(u0, p) || continue
+                __compatible(u0, alg) || continue
+                __compatible(u0, Val(mode)) || continue
+                __compatible(alg, Val(mode)) || continue
 
-            sol = solve(NonlinearProblem(test_f, u0, p), alg)
-            if SciMLBase.successful_retcode(sol)
-                gs = abs.(ForwardDiff.jacobian(solve_with(Val{mode}(), u0, alg), p))
-                gs_true = abs.(jacobian_f(u0, p))
-                if !(isapprox(gs, gs_true, atol = 1e-5))
-                    @show sol.retcode, sol.u
-                    @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_jacobian=gs true_jacobian=gs_true
-                else
-                    @test abs.(gs)≈abs.(gs_true) atol=1e-5
+                sol = solve(NonlinearProblem(test_f, u0, p), alg)
+                if SciMLBase.successful_retcode(sol)
+                    gs = abs.(ForwardDiff.jacobian(solve_with(Val{mode}(), u0, alg), p))
+                    gs_true = abs.(jacobian_f(u0, p))
+                    if !(isapprox(gs, gs_true, atol = 1e-5))
+                        @show sol.retcode, sol.u
+                        @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_jacobian=gs true_jacobian=gs_true
+                    else
+                        @test abs.(gs)≈abs.(gs_true) atol=1e-5
+                    end
                 end
             end
         end

test/core/nlls.jl renamed to test/core/nlls_tests.jl

@@ -1,5 +1,7 @@
-using NonlinearSolve,
-    LinearSolve, LinearAlgebra, XUnit, StableRNGs, Random, ForwardDiff, Zygote
+@testsetup module CoreNLLSTesting
+using Reexport
+@reexport using NonlinearSolve,
+    LinearSolve, LinearAlgebra, StableRNGs, Random, ForwardDiff, Zygote
 
 true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])
 true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]))
@@ -22,11 +24,6 @@ function loss_function(resid, θ, p)
 end
 
 θ_init = θ_true .+ randn!(StableRNG(0), similar(θ_true)) * 0.1
-prob_oop = NonlinearLeastSquaresProblem{false}(loss_function, θ_init, x)
-prob_iip = NonlinearLeastSquaresProblem(NonlinearFunction(loss_function;
-        resid_prototype = zero(y_target)), θ_init, x)
-
-nlls_problems = [prob_oop, prob_iip]
 
 solvers = []
 for linsolve in [nothing, LUFactorization(), KrylovJL_GMRES(), KrylovJL_LSMR()]
@@ -52,36 +49,45 @@ for radius_update_scheme in [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.Noc
     push!(solvers, TrustRegion(; radius_update_scheme))
 end
 
-@testcase "General NLLS Solvers" begin
+export solvers, θ_init, x, y_target, true_function, θ_true, loss_function
+end
+
+@testitem "General NLLS Solvers" setup=[CoreNLLSTesting] begin
+    prob_oop = NonlinearLeastSquaresProblem{false}(loss_function, θ_init, x)
+    prob_iip = NonlinearLeastSquaresProblem(NonlinearFunction(loss_function;
+            resid_prototype = zero(y_target)), θ_init, x)
+
+    nlls_problems = [prob_oop, prob_iip]
+
     for prob in nlls_problems, solver in solvers
         sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)
         @test SciMLBase.successful_retcode(sol)
         @test maximum(abs, sol.resid) < 1e-6
     end
 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
+@testitem "Custom VJP" setup=[CoreNLLSTesting] begin
+    # 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), transpose(J), v)
-    return nothing
-end
+    function vjp!(Jv, v, θ, p)
+        resid = zeros(length(p))
+        J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ)
+        mul!(vec(Jv), transpose(J), v)
+        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),
-]
+    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),
+    ]
 
-@testcase "Custom VJP" begin
     for prob in probs, solver in solvers
         sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8)
         @test maximum(abs, sol.resid) < 1e-6