Start implementing trimming test

We introduce `SafeTestsets` as part of this, inspired by the way the
downstream tests are set up in the `SciMLBase` repo.

Author: RomeoV

Project.toml

@@ -143,6 +143,7 @@ PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -170,5 +171,4 @@ path = "lib/NonlinearSolveSpectralMethods"
 path = "lib/SimpleNonlinearSolve"
 
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"]
-
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SafeTestsets", "SIAMFANLEquations", "SparseConnectivityTracer", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"]

test/runtests.jl

@@ -1,9 +1,18 @@
-using ReTestItems, NonlinearSolve, Hwloc, InteractiveUtils, Pkg
+using NonlinearSolve, Hwloc, InteractiveUtils, Pkg
+using SafeTestsets
+using ReTestItems
 
 @info sprint(InteractiveUtils.versioninfo)
 
 const GROUP = lowercase(get(ENV, "GROUP", "All"))
 
+function activate_trim_env!()
+    Pkg.activate(abspath(joinpath(dirname(@__FILE__), "trim")))
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
+    return nothing
+end
+
 const EXTRA_PKGS = Pkg.PackageSpec[]
 if GROUP == "all" || GROUP == "downstream"
     push!(EXTRA_PKGS, Pkg.PackageSpec("ModelingToolkit"))
@@ -12,11 +21,13 @@ end
 length(EXTRA_PKGS) ≥ 1 && Pkg.add(EXTRA_PKGS)
 
 const RETESTITEMS_NWORKERS = parse(
-    Int, get(ENV, "RETESTITEMS_NWORKERS",
+    Int, get(
+        ENV, "RETESTITEMS_NWORKERS",
         string(min(ifelse(Sys.iswindows(), 0, Hwloc.num_physical_cores()), 4))
     )
 )
-const RETESTITEMS_NWORKER_THREADS = parse(Int,
+const RETESTITEMS_NWORKER_THREADS = parse(
+    Int,
     get(
         ENV, "RETESTITEMS_NWORKER_THREADS",
         string(max(Hwloc.num_virtual_cores() ÷ max(RETESTITEMS_NWORKERS, 1), 1))
@@ -25,8 +36,22 @@ const RETESTITEMS_NWORKER_THREADS = parse(Int,
 
 @info "Running tests for group: $(GROUP) with $(RETESTITEMS_NWORKERS) workers"
 
-ReTestItems.runtests(
-    NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]),
-    nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS,
-    testitem_timeout = 3600
-)
+if GROUP != "trim"
+    ReTestItems.runtests(
+        NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]),
+        nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS,
+        testitem_timeout = 3600
+    )
+elseif GROUP == "trim" && VERSION >= v"1.12.0-rc1"  # trimming has been introduced in julia 1.12
+    activate_trim_env!()
+    @safetestset "Clean implementation (non-trimmable)" begin
+        using SciMLBase: successful_retcode
+        include("trim/clean_optimization.jl")
+        @test successful_retcode(minimize(1.0).retcode)
+    end
+    @safetestset "Trimmable implementation" begin
+        using SciMLBase: successful_retcode
+        include("trim/trimmable_optimization.jl")
+        @test successful_retcode(minimize(1.0).retcode)
+    end
+end

test/trim/Project.toml

@@ -0,0 +1,27 @@
+[deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
+PolyesterWeave = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+
+[sources]
+ForwardDiff = {rev = "rv/remove-quote-assert-string-interpolation", url = "https://github.com/RomeoV/ForwardDiff.jl"}
+LinearSolve = {rev = "rv/remove-linsolve-forwarddiff-special-path", url = "https://github.com/RomeoV/LinearSolve.jl"}
+NonlinearSolveFirstOrder = {path = "../../lib/NonlinearSolveFirstOrder"}
+Polyester = {rev = "master", url = "https://github.com/RomeoV/Polyester.jl"}
+PolyesterWeave = {rev = "main", url = "https://github.com/RomeoV/PolyesterWeave.jl"}
+SciMLBase = {rev = "as/fix-jet-opt", url = "https://github.com/AayushSabharwal/SciMLBase.jl"}
+
+[compat]
+ADTypes = "1.15.0"
+DiffEqBase = "6.179.0"
+LinearAlgebra = "1.12.0"
+NonlinearSolveFirstOrder = "1.6.0"
+StaticArrays = "1.9.0"

test/trim/clean_optimization.jl

@@ -0,0 +1,32 @@
+using NonlinearSolveFirstOrder
+using ADTypes: AutoForwardDiff
+using ForwardDiff
+using LinearAlgebra
+using StaticArrays
+using LinearSolve
+const LS = LinearSolve
+
+function f(u, p)
+    L, U = cholesky(p.Σ)
+    rhs = (u .* u .- p.λ)
+    # there are some issues currently with LinearSolve and triangular matrices,
+    # so we just make `L` dense here.
+    linprob = LinearProblem(Matrix(L), rhs)
+    alg = LS.GenericLUFactorization()
+    sol = LinearSolve.solve(linprob, alg)
+    return sol.u
+end
+
+struct MyParams{T, M}
+    λ::T
+    Σ::M
+end
+
+function minimize(x)
+    autodiff = AutoForwardDiff(; chunksize=1)
+    alg = TrustRegion(; autodiff, linsolve=LS.CholeskyFactorization())
+    ps = MyParams(rand(), hermitianpart(rand(2,2)+2I))
+    prob = NonlinearLeastSquaresProblem{false}(f, rand(2), ps)
+    sol = solve(prob, alg)
+    return sol
+end

test/trim/trimmable_optimization.jl

@@ -0,0 +1,35 @@
+using NonlinearSolveFirstOrder
+using ADTypes: AutoForwardDiff
+using ForwardDiff
+using LinearAlgebra
+using StaticArrays
+using LinearSolve
+const LS = LinearSolve
+
+function f(u, p)
+    L, U = cholesky(p.Σ)
+    rhs = (u .* u .- p.λ)
+    # there are some issues currently with LinearSolve and triangular matrices,
+    # so we just make `L` dense here.
+    linprob = LinearProblem(Matrix(L), rhs)
+    alg = LS.GenericLUFactorization()
+    sol = LinearSolve.solve(linprob, alg)
+    return sol.u
+end
+
+struct MyParams{T, M}
+    λ::T
+    Σ::M
+end
+
+const autodiff = AutoForwardDiff(; chunksize = 1)
+const alg = TrustRegion(; autodiff, linsolve = LS.CholeskyFactorization())
+const prob = NonlinearLeastSquaresProblem{false}(f, rand(2), MyParams(rand(), hermitianpart(rand(2, 2) + 2I)))
+const cache = init(prob, alg)
+
+function minimize(x)
+    ps = MyParams(x, hermitianpart(rand(2, 2) + 2I))
+    reinit!(cache, rand(2); p = ps)
+    solve!(cache)
+    return cache
+end