Add precompile for NLLS

SciML · Nov 2, 2023 · 0d89100 · 0d89100
1 parent 40a9ded
commit 0d89100
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 13 deletions.
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -5,9 +5,9 @@ if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@max_m
 end
 
 import Reexport: @reexport
-import PrecompileTools
+import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workload
 
-PrecompileTools.@recompile_invalidations begin
+@recompile_invalidations begin
     using DiffEqBase, LinearAlgebra, LinearSolve, SparseArrays, SparseDiffTools
     using FastBroadcast: @..
     import ArrayInterface: restructure
@@ -95,17 +95,49 @@ include("jacobian.jl")
 include("ad.jl")
 include("default.jl")
 
-PrecompileTools.@compile_workload begin
-    for T in (Float32, Float64)
-        probs = (NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2)),
-            NonlinearProblem{false}((u, p) -> u .* u .- p, T[0.1], T[2]),
-            NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p, T[0.1], T[2]))
+@setup_workload begin
+    nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1),
+        (NonlinearFunction{false}((u, p) -> u .* u .- p), [0.1]),
+        (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1]))
+    probs_nls = NonlinearProblem[]
+    for T in (Float32, Float64), (fn, u0) in nlfuncs
+        push!(probs_nls, NonlinearProblem(fn, T.(u0), T(2)))
+    end
 
-        precompile_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(),
-            PseudoTransient(), GeneralBroyden(), GeneralKlement(), DFSane(), nothing)
+    nls_algs = (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), PseudoTransient(),
+        GeneralBroyden(), GeneralKlement(), DFSane(), nothing)
+
+    probs_nlls = NonlinearLeastSquaresProblem[]
+    nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]),
+        (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]),
+        (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p,
+                resid_prototype = zeros(1)), [0.1, 0.0]),
+        (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
+                resid_prototype = zeros(4)), [0.1, 0.1]))
+    for (fn, u0) in nlfuncs
+        push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0))
+    end
+    nlfuncs = ((NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), Float32[0.1, 0.0]),
+        (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)),
+            Float32[0.1, 0.1]),
+        (NonlinearFunction{true}((du, u, p) -> du[1] = u[1] * u[1] - p,
+                resid_prototype = zeros(Float32, 1)), Float32[0.1, 0.0]),
+        (NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p),
+                resid_prototype = zeros(Float32, 4)), Float32[0.1, 0.1]))
+    for (fn, u0) in nlfuncs
+        push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0f0))
+    end
 
-        for prob in probs, alg in precompile_algs
-            solve(prob, alg, abstol = T(1e-2))
+    nlls_algs = (LevenbergMarquardt(), GaussNewton(),
+        LevenbergMarquardt(; linsolve = LUFactorization()),
+        GaussNewton(; linsolve = LUFactorization()))
+
+    @compile_workload begin
+        for prob in probs_nls, alg in nls_algs
+            solve(prob, alg, abstol = 1e-2)
+        end
+        for prob in probs_nlls, alg in nlls_algs
+            solve(prob, alg, abstol = 1e-2)
         end
     end
 end

diff --git a/test/basictests.jl b/test/basictests.jl
@@ -390,7 +390,7 @@ end
     @testset "Keyword Arguments" begin
         damping_initial = [0.5, 2.0, 5.0]
         damping_increase_factor = [1.5, 3.0, 10.0]
-        damping_decrease_factor = Float64[2, 5, 12]
+        damping_decrease_factor = Float64[2, 5, 10.0]
         finite_diff_step_geodesic = [0.02, 0.2, 0.3]
         α_geodesic = [0.6, 0.8, 0.9]
         b_uphill = Float64[0, 1, 2]
@@ -408,7 +408,7 @@ end
                 b_uphill = options[6], min_damping_D = options[7])
 
             probN = NonlinearProblem{false}(quadratic_f, [1.0, 1.0], 2.0)
-            sol = solve(probN, alg, abstol = 1e-12)
+            sol = solve(probN, alg; abstol = 1e-13, maxiters = 10000)
             @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-10)
         end
     end