lotka_sensitivity.jl

using ComponentArrays, DifferentialEquations, Plots, UnPack


## Functions
function lotka!(D, x, p, t)
    @unpack NB, NR = x
    @unpack ϵ₁, ϵ₂, γ₁, γ₂ = p

    D.NB =  ϵ₁*NB - γ₁*NR*NB
    D.NR = -ϵ₂*NR + γ₂*NR*NB
    return nothing
end

function param_jac!(J, x, p)
    @unpack NB, NR = x
    @unpack ϵ₁, ϵ₂, γ₁, γ₂ = p

    J[Val(:NB), Val(:ϵ₁)] =  NB
    J[Val(:NB), Val(:γ₁)] = -NR*NB
    J[Val(:NR), Val(:ϵ₂)] = -NR
    J[Val(:NR), Val(:γ₂)] =  NR*NB
    return J
end

function state_jac!(J, x, p)
    @unpack NB, NR = x
    @unpack ϵ₁, ϵ₂, γ₁, γ₂ = p

    J[Val(:NB), Val(:NB)] =  ϵ₁ - γ₁*NR
    J[Val(:NB), Val(:NR)] = -γ₁*NB
    J[Val(:NR), Val(:NB)] =  γ₂*NR
    J[Val(:NR), Val(:NR)] = -ϵ₂ + γ₂*NB
    return J
end

function lotka_sensitivity!(D, vars, p, t)
    @unpack x, s = vars
    @unpack params, jacobians = p
    Jp, Jx = jacobians.param, jacobians.state

    param_jac!(Jp, x, params)
    state_jac!(Jx, x, params)

    lotka!(D.x, x, params, t)
    
    D.s .= Jx*s .+ Jp
    return nothing
end

## Inputs
x0 = ComponentArray(NB=1.0, NR=1.0)
params = ComponentArray(ϵ₁=3.0, ϵ₂=1.0, γ₁=0.5, γ₂=0.4)

s0 = x0 * params' * 0
jacobians = (
    state = x0 * x0' * 0,
    param = x0 * params' * 0,
)

u0 = ComponentArray(x=x0, s=s0)
p = (params=params, jacobians=jacobians)

t = (0.0, 3.0)


## Solve
prob = ODEProblem(lotka!, x0, t, params)
sol = solve(prob)

prob_sens = ODEProblem(lotka_sensitivity!, u0, t, p)
sol_sens = solve(prob_sens)


## Plot
p1 = plot(sol, vars=1:2)
plot!(sol_sens, vars=1:2, linestyle=:dash)

p2 = plot(sol_sens)

plot(p1, p2, layout=(2,1), size=(800,600))