Use batching for gsa

README.md

@@ -6,7 +6,7 @@
 [![codecov](https://codecov.io/gh/SciML/EasyModelAnalysis.jl/branch/main/graph/badge.svg)](https://app.codecov.io/gh/SciML/EasyModelAnalysis.jl)
 [![Build Status](https://github.com/SciML/EasyModelAnalysis.jl/workflows/CI/badge.svg)](https://github.com/SciML/EasyModelAnalysis.jl/actions?query=workflow%3ACI)
 
-[![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
+[![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor%27s%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
 [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle)
 
 EasyModelAnalysis does exactly what it says: it makes model analysis easy. Want to know the first time
@@ -17,8 +17,8 @@ one-liner queries over SciML-defined differential equation models.
 
 ## Tutorials and Documentation
 
-For information on using the package, see the [stable documentation](https://docs.sciml.ai/EasyModelAnalysis/stable/). 
-Use the [in-development documentation](https://docs.sciml.ai/EasyModelAnalysis/dev/) for the version of the documentation 
+For information on using the package, see the [stable documentation](https://docs.sciml.ai/EasyModelAnalysis/stable/).
+Use the [in-development documentation](https://docs.sciml.ai/EasyModelAnalysis/dev/) for the version of the documentation
 which contains the unreleased features.
 
 ## Quick Demonstration

docs/src/tutorials/sensitivity_analysis.md

@@ -45,7 +45,7 @@ Thus for example, let's calculate the sensitivity of `y(100)` over the parameter
 
 ```@example sensitivity
 pbounds = [ρ => [0.0, 20.0], β => [0.0, 100.0]]
-sensres = get_sensitivity(prob, 100.0, y, pbounds; samples = 10000)
+sensres = get_sensitivity(prob, 100.0, y, pbounds)
 ```
 
 The output shows values of `first_order`, `second_order` and `total_order` sensitivities. These are quantities that define the

docs/src/tutorials/threshold_interventions.md

@@ -59,8 +59,8 @@ opt_tspan, (s1, s2, s3), ret = optimal_threshold_intervention(prob, [p => -1.0],
 opt_tspan
 ```
 
-The `opt_tspan` gives us the optimal timespan of the intervention: we should begin the decimation of bunny civilization when it reaches 
-5.15 months, and then we can turn the destruction device off at 27.8 months into our tourism operation. To see the effect of this we 
+The `opt_tspan` gives us the optimal timespan of the intervention: we should begin the decimation of bunny civilization when it reaches
+5.15 months, and then we can turn the destruction device off at 27.8 months into our tourism operation. To see the effect of this we
 can plot the results of the three sub-intervals:
 
 ```@example threshold_intervention

src/sensitivity.jl

@@ -1,20 +1,27 @@
 function _get_sensitivity(prob, t, x, pbounds; samples)
     boundvals = getfield.(pbounds, :second)
     boundkeys = getfield.(pbounds, :first)
-    function f(p)
-        prob = remake(prob; p = Pair.(boundkeys, p))
-        sol = solve(prob, saveat = t)
-        sol(t; idxs = x)
+    f = function (p)
+        prob_func(prob, i, repeat) = remake(prob; p = Pair.(boundkeys, p[:, i]))
+        ensemble_prob = EnsembleProblem(prob, prob_func = prob_func)
+        sol = solve(ensemble_prob, Tsit5(), EnsembleThreads(); saveat = t,
+                    trajectories = size(p, 2))
+        out = zeros(size(p, 2))
+        for i in 1:size(p, 2)
+            out[i] = sol[i](t; idxs = x)
+        end
+        out
     end
-    return GlobalSensitivity.gsa(f, Sobol(; order = [0, 1, 2]), boundvals; samples)
+    return GlobalSensitivity.gsa(f, Sobol(; order = [0, 1, 2]), boundvals; samples,
+                                 batch = true)
 end
 
 """
     get_sensitivity(prob, t, x, pbounds)
 
 Returns the sensitivity of the solution at time `t` and state `x` to the parameters in `pbounds`.
 """
-function get_sensitivity(prob, t, x, pbounds; samples = 1000)
+function get_sensitivity(prob, t, x, pbounds; samples = 10000)
     sensres = _get_sensitivity(prob, t, x, pbounds; samples)
     boundvals = getfield.(pbounds, :second)
     boundkeys = getfield.(pbounds, :first)
@@ -24,7 +31,7 @@ function get_sensitivity(prob, t, x, pbounds; samples = 1000)
         res_dict[Symbol(boundkeys[i], "_total_order")] = sensres.ST[i]
     end
     for i in eachindex(boundkeys)
-        for j in i:length(boundkeys)
+        for j in (i + 1):length(boundkeys)
             res_dict[Symbol(boundkeys[i], "_", boundkeys[j], "_second_order")] = sensres.S2[i,
                                                                                             j]
         end