Global Sensitivity Analysis of the Lotka-Volterra model

The tutorial covers a workflow of using GlobalSensitivity.jl on the Lotka-Volterra differential equations. We showcase how to use multiple GSA methods, analyze their results and leverage Julia's parallelism capabilities to perform Global Sensitivity analysis at scale.

using GlobalSensitivity, QuasiMonteCarlo, OrdinaryDiffEq, Statistics, CairoMakie

function f(du, u, p, t)
    du[1] = p[1] * u[1] - p[2] * u[1] * u[2] #prey
    du[2] = -p[3] * u[2] + p[4] * u[1] * u[2] #predator
end

u0 = [1.0; 1.0]
tspan = (0.0, 10.0)
p = [1.5, 1.0, 3.0, 1.0]
prob = ODEProblem(f, u0, tspan, p)
t = collect(range(0, stop = 10, length = 200))

f1 = function (p)
    prob1 = remake(prob; p = p)
    sol = solve(prob1, Tsit5(); saveat = t)
    return [mean(sol[1, :]), maximum(sol[2, :])]
end

bounds = [[1, 5], [1, 5], [1, 5], [1, 5]]

reg_sens = gsa(f1, RegressionGSA(true), bounds, samples = 200)
fig = Figure(resolution = (600, 400))
ax, hm = CairoMakie.heatmap(fig[1, 1],
    reg_sens.partial_correlation,
    axis = (xticksvisible = false, yticksvisible = false, yticklabelsvisible = false,
        xticklabelsvisible = false, title = "Partial correlation"))
Colorbar(fig[1, 2], hm)
ax, hm = CairoMakie.heatmap(fig[2, 1],
    reg_sens.standard_regression,
    axis = (xticksvisible = false, yticksvisible = false, yticklabelsvisible = false,
        xticklabelsvisible = false, title = "Standard regression"))
Colorbar(fig[2, 2], hm)
fig

using StableRNGs
_rng = StableRNG(1234)
morris_sens = gsa(f1, Morris(), bounds, rng = _rng)
fig = Figure(resolution = (600, 400))
CairoMakie.scatter(fig[1, 1], [1, 2, 3, 4], morris_sens.means_star[1, :], color = :green,
    axis = (xticksvisible = false, xticklabelsvisible = false, title = "Prey"))
CairoMakie.scatter(fig[1, 2], [1, 2, 3, 4], morris_sens.means_star[2, :], color = :red,
    axis = (xticksvisible = false, xticklabelsvisible = false, title = "Predator"))
fig

sobol_sens = gsa(f1, Sobol(), bounds, samples = 500)
efast_sens = gsa(f1, eFAST(), bounds, samples = 500)

GlobalSensitivity.eFASTResult{Matrix{Float64}}([0.0007338213524505229 0.0005459213694739437 0.3563481753250383 0.5463286506565318; 0.17893455695758348 0.527353927368662 0.005431762608263231 0.011403252063478406], [0.015479300016764008 0.014913630705892023 0.4503957662541992 0.6429548637211346; 0.36649135782530984 0.7604727856981315 0.055201167489284275 0.07052152581469417])

fig = Figure(resolution = (600, 400))
barplot(fig[1, 1],
    [1, 2, 3, 4],
    sobol_sens.S1[1, :],
    color = :green,
    axis = (xticksvisible = false, xticklabelsvisible = false,
        title = "Prey (Sobol)", ylabel = "First order"))
barplot(fig[2, 1], [1, 2, 3, 4], sobol_sens.ST[1, :], color = :green,
    axis = (xticksvisible = false, xticklabelsvisible = false, ylabel = "Total order"))
barplot(fig[1, 2], [1, 2, 3, 4], efast_sens.S1[1, :], color = :red,
    axis = (xticksvisible = false, xticklabelsvisible = false, title = "Prey (eFAST)"))
barplot(fig[2, 2], [1, 2, 3, 4], efast_sens.ST[1, :], color = :red,
    axis = (xticksvisible = false, xticklabelsvisible = false))
fig

fig = Figure(resolution = (600, 400))
barplot(fig[1, 1],
    [1, 2, 3, 4],
    sobol_sens.S1[2, :],
    color = :green,
    axis = (xticksvisible = false, xticklabelsvisible = false,
        title = "Predator (Sobol)", ylabel = "First order"))
barplot(fig[2, 1], [1, 2, 3, 4], sobol_sens.ST[2, :], color = :green,
    axis = (xticksvisible = false, xticklabelsvisible = false, ylabel = "Total order"))
barplot(fig[1, 2], [1, 2, 3, 4], efast_sens.S1[2, :], color = :red,
    axis = (xticksvisible = false, xticklabelsvisible = false, title = "Predator (eFAST)"))
barplot(fig[2, 2], [1, 2, 3, 4], efast_sens.ST[2, :], color = :red,
    axis = (xticksvisible = false, xticklabelsvisible = false))
fig

using QuasiMonteCarlo
samples = 500
lb = [1.0, 1.0, 1.0, 1.0]
ub = [5.0, 5.0, 5.0, 5.0]
sampler = SobolSample()
A, B = QuasiMonteCarlo.generate_design_matrices(samples, lb, ub, sampler)
sobol_sens_desmat = gsa(f1, Sobol(), A, B)

f_batch = function (p)
    prob_func(prob, i, repeat) = remake(prob; p = p[:, i])
    ensemble_prob = EnsembleProblem(prob, prob_func = prob_func)

    sol = solve(
        ensemble_prob, Tsit5(), EnsembleThreads(); saveat = t, trajectories = size(p, 2))

    out = zeros(2, size(p, 2))

    for i in 1:size(p, 2)
        out[1, i] = mean(sol[i][1, :])
        out[2, i] = maximum(sol[i][2, :])
    end

    return out
end

sobol_sens_batch = gsa(f_batch, Sobol(), A, B, batch = true)

@time gsa(f1, Sobol(), A, B)
@time gsa(f_batch, Sobol(), A, B, batch = true)

GlobalSensitivity.SobolResult{Matrix{Float64}, Nothing, Nothing, Nothing}([0.00326173894105304 0.000847340115491134 0.3861155616639945 0.5857005480643274; 0.18030754948552288 0.6811024027388886 0.015497253502704226 0.01678762044781814], nothing, nothing, nothing, [0.0030820547740310414 0.0017822314786998207 0.5227660383161423 0.6434141725960043; 0.3794747144680316 0.8342675343550613 0.052815796905431804 0.0522308865963962], nothing)

f1 = function (p)
    prob1 = remake(prob; p = p)
    sol = solve(prob1, Tsit5(); saveat = t)
end
sobol_sens = gsa(f1, Sobol(nboot = 20), bounds, samples = 500)
fig = Figure(resolution = (600, 400))
ax, hm = CairoMakie.scatter(
    fig[1, 1], sobol_sens.S1[1][1, 2:end], label = "Prey", markersize = 4)
CairoMakie.scatter!(
    fig[1, 1], sobol_sens.S1[1][2, 2:end], label = "Predator", markersize = 4)

# Legend(fig[1,2], ax)

ax, hm = CairoMakie.scatter(
    fig[1, 2], sobol_sens.S1[2][1, 2:end], label = "Prey", markersize = 4)
CairoMakie.scatter!(
    fig[1, 2], sobol_sens.S1[2][2, 2:end], label = "Predator", markersize = 4)

ax, hm = CairoMakie.scatter(
    fig[2, 1], sobol_sens.S1[3][1, 2:end], label = "Prey", markersize = 4)
CairoMakie.scatter!(
    fig[2, 1], sobol_sens.S1[3][2, 2:end], label = "Predator", markersize = 4)

ax, hm = CairoMakie.scatter(
    fig[2, 2], sobol_sens.S1[4][1, 2:end], label = "Prey", markersize = 4)
CairoMakie.scatter!(
    fig[2, 2], sobol_sens.S1[4][2, 2:end], label = "Predator", markersize = 4)

title = Label(fig[0, :], "First order Sobol indices")
legend = Legend(fig[2, 3], ax)
fig