Precompile NLLS Workflows as well

SciML · Nov 2, 2023 · 90d1674 · 90d1674
1 parent 5011657
commit 90d1674
Showing 1 changed file with 64 additions and 0 deletions.
diff --git a/ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl b/ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl
@@ -55,6 +55,70 @@ end
             solve(prob, alg)
         end
     end
+
+    function f1_nlls!(du, u, p, t)
+        du[1] = u[2]
+        du[2] = -u[1]
+    end
+
+    f1_nlls(u, p, t) = [u[2], -u[1]]
+
+    function bc1_nlls!(resid, sol, p, t)
+        solₜ₁ = sol(0.0)
+        solₜ₂ = sol(100.0)
+        resid[1] = solₜ₁[1]
+        resid[2] = solₜ₂[1] - 1
+        resid[3] = solₜ₂[2] + 1.729109
+        return nothing
+    end
+    bc1_nlls(sol, p, t) = [sol(0.0)[1], sol(100.0)[1] - 1, sol(100.0)[2] + 1.729109]
+
+    bc1_nlls_a!(resid, ua, p) = (resid[1] = ua[1])
+    bc1_nlls_b!(resid, ub, p) = (resid[1] = ub[1] - 1; resid[2] = ub[2] + 1.729109)
+
+    bc1_nlls_a(ua, p) = [ua[1]]
+    bc1_nlls_b(ub, p) = [ub[1] - 1, ub[2] + 1.729109]
+
+    tspan = (0.0, 100.0)
+    u0 = [0.0, 1.0]
+    bcresid_prototype1 = Array{Float64}(undef, 3)
+    bcresid_prototype2 = (Array{Float64}(undef, 1), Array{Float64}(undef, 2))
+
+    probs = [
+        BVProblem(BVPFunction(f1_nlls!, bc1_nlls!; bcresid_prototype = bcresid_prototype1),
+            u0, tspan),
+        BVProblem(BVPFunction(f1_nlls, bc1_nlls; bcresid_prototype = bcresid_prototype1),
+            u0, tspan),
+        TwoPointBVProblem(f1_nlls!, (bc1_nlls_a!, bc1_nlls_b!), u0, tspan;
+            bcresid_prototype = bcresid_prototype2),
+        TwoPointBVProblem(f1_nlls, (bc1_nlls_a, bc1_nlls_b), u0, tspan;
+            bcresid_prototype = bcresid_prototype2),
+    ]
+
+    algs = [
+        Shooting(Tsit5();
+            nlsolve = LevenbergMarquardt(; autodiff = AutoForwardDiff(chunksize = 2))),
+        Shooting(Tsit5();
+            nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2))),
+        MultipleShooting(10,
+            Tsit5();
+            nlsolve = LevenbergMarquardt(; autodiff = AutoForwardDiff(chunksize = 2)),
+            jac_alg = BVPJacobianAlgorithm(;
+                bc_diffmode = AutoForwardDiff(; chunksize = 2),
+                nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
+        MultipleShooting(10,
+            Tsit5();
+            nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2)),
+            jac_alg = BVPJacobianAlgorithm(;
+                bc_diffmode = AutoForwardDiff(; chunksize = 2),
+                nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
+    ]
+
+    @compile_workload begin
+        for prob in probs, alg in algs
+            solve(prob, alg)
+        end
+    end
 end
 
 end