Use preferences

SciML · Nov 3, 2023 · 347f22a · 347f22a
1 parent 33f9763
commit 347f22a
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 65 deletions.
diff --git a/README.md b/README.md
@@ -48,6 +48,9 @@ Precompilation can be controlled via `Preferences.jl`
 - `PrecompileMIRK` -- Precompile the MIRK2 - MIRK6 algorithms (default: `true`).
 - `PrecompileShooting` -- Precompile the single shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
 - `PrecompileMultipleShooting` -- Precompile the multiple shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
+- `PrecompileMIRKNLLS` -- Precompile the MIRK2 - MIRK6 algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above).
+- `PrecompileShootingNLLS` -- Precompile the single shooting algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above). This is triggered when `OrdinaryDiffEq` is loaded.
+- `PrecompileMultipleShootingNLLS` -- Precompile the multiple shooting algorithms for under-determined and over-determined BVPs (default: `true` on Julia Version 1.10 and above). This is triggered when `OrdinaryDiffEq` is loaded.
 
 To set these preferences before loading the package, do the following (replacing `PrecompileShooting` with the preference you want to set, or pass in multiple pairs to set them together):
 

diff --git a/ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl b/ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl
@@ -64,69 +64,78 @@ end
         end
     end
 
-    if VERSION ≥ v"1.10-"
-        function f1_nlls!(du, u, p, t)
-            du[1] = u[2]
-            du[2] = -u[1]
-        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]]
+    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))),
-        ]
+    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]
 
-        @compile_workload begin
-            for prob in probs, alg in algs
-                solve(prob, alg)
-            end
+    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 = []
+
+    if @load_preference("PrecompileShootingNLLS", VERSION≥v"1.10-")
+        append!(algs,
+            [
+                Shooting(Tsit5();
+                    nlsolve = LevenbergMarquardt(;
+                        autodiff = AutoForwardDiff(chunksize = 2))),
+                Shooting(Tsit5();
+                    nlsolve = GaussNewton(; autodiff = AutoForwardDiff(chunksize = 2))),
+            ])
+    end
+
+    if @load_preference("PrecompileMultipleShootingNLLS", VERSION≥v"1.10-")
+        append!(algs,
+            [
+                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))),
+            ])
+    end
+
+    @compile_workload begin
+        for prob in probs, alg in algs
+            solve(prob, alg)
         end
     end
 end

diff --git a/src/BoundaryValueDiffEq.jl b/src/BoundaryValueDiffEq.jl
@@ -133,12 +133,21 @@ end
 
     nlsolvers = [LevenbergMarquardt(), GaussNewton()]
 
-    @compile_workload begin
-        for prob in probs, nlsolve in nlsolvers,
-            alg in (MIRK2(; jac_alg, nlsolve), MIRK3(; jac_alg, nlsolve),
-                MIRK4(; jac_alg, nlsolve), MIRK5(; jac_alg, nlsolve),
-                MIRK6(; jac_alg, nlsolve))
+    algs = []
 
+    if Preferences.@load_preference("PrecompileMIRKNLLS", VERSION≥v"1.10-")
+        for nlsolve in nlsolvers
+            append!(algs,
+                [
+                    MIRK2(; jac_alg, nlsolve), MIRK3(; jac_alg, nlsolve),
+                    MIRK4(; jac_alg, nlsolve), MIRK5(; jac_alg, nlsolve),
+                    MIRK6(; jac_alg, nlsolve),
+                ])
+        end
+    end
+
+    @compile_workload begin
+        for prob in probs, alg in algs
             solve(prob, alg; dt = 0.2)
         end
     end