Direct Forward Sensitivity Analysis of ODEs in ModelingToolkit.jl for a Subset of Parameters

Hello,

I am working on a project that involves setting up and solving an ODE system. I’ve defined my ODE system using ModelingToolkit.jl and not as a function. My ODE system has many parameters, and I’m interested in performing a Direct Forward Sensitivity Analysis on a subset of these parameters, not all of them. Post-analysis, I intend to apply Bayesian inference to the results further.

Here’s a simplified overview of what I am trying to do:

using ModelingToolkit, Symbolics, DifferentialEquations, SciMLSensitivity

@variables t, x(t), y(t)
@parameters p1, p2, p3, p4

D = Differential(t)

# ODE system definition
eqs = [D(x) ~ p1 * x + p2 * y,
       D(y) ~ p3 * y + p4 * x]

# A subset of parameters I want to consider for sensitivity analysis
param_subset = [p1, p3]

In my case, the number of parameters is vast, and I have a stiff problem. I want to construct a sensitivity matrix from them and then analyze it. How can I construct ODEForwardSensitivityProblem for a subset of parameters?

Thank you in advance,

@ArnoStrouwen how hard would it be to revive https://github.com/SciML/ModelingToolkit.jl/pull/398 to merge?

I tried updated the example in that PR to modern MTK syntax.

using ModelingToolkit
using OrdinaryDiffEq
@parameters t #time
D = Differential(t)

@variables Cₛ(t), Cₓ(t), ν(t) # states
x = [Cₛ, Cₓ, ν]
n_x = length(x)
@parameters uᵢₙ # control input
u = [uᵢₙ]
n_u = length(u)
@parameters μₘₐₓ Cₛ_ᵢₙ Cₛ₀ Cₓ₀ # uncertain parameters
p = [μₘₐₓ, Cₛ_ᵢₙ, Cₛ₀, Cₓ₀]
n_p = length(p)
x₀ = [Cₛ₀, Cₓ₀, μₘₐₓ] # parametrized initial conditions; note third initial condition is also present in the dynamics
@parameters  Kₛ, yₓ_ₛ, m #other constants
cons = [Kₛ, yₓ_ₛ, m]
μ = μₘₐₓ*Cₛ/Kₛ + Cₛ
σ = μ/yₓ_ₛ + m
all_parameters = [p;cons;u]
eqs = [D(Cₛ) ~ -σ*Cₓ + uᵢₙ/ν*Cₛ_ᵢₙ - uᵢₙ/ν*Cₛ,
       D(Cₓ) ~ μ*Cₓ - uᵢₙ/ν*Cₓ,
       D(ν) ~ uᵢₙ]
@named sys = ODESystem(eqs)
tspan =  (0.0, 10.0)
prob = ODEProblem(sys, x.=>x₀, tspan, all_parameters .=> rand(length(all_parameters)))
solve(prob,Tsit5())

function forward_sensitivity_transform(sys,p)
    iv  = sys.iv
    f = du_dt = [eq.rhs for eq ∈ sys.eqs]
    f = collect(f)
    u = states(sys)
    D_p = [Differential(par) for par ∈ p]
    sen_symbols = [Symbol(:d,state.metadata.value[2],:_d,par.val.metadata.parent.value[2]) for state ∈ u, par in p]
    du_dp = [(@variables $sen_symbol(t))[1] for sen_symbol ∈ sen_symbols]
    D = Differential(iv)
    uₑₓₜ= vcat(u,vec(du_dp))
    df_du = ModelingToolkit.jacobian(f,u)
    df_dp = ModelingToolkit.jacobian(f,p)
    ddu_dpdt = df_du*du_dp + df_dp
    duₑₓₜ_dt = fₑₓₜ = vcat(du_dt,vec(ddu_dpdt))
    eqsₑₓₜ = simplify.(D.(uₑₓₜ) .~ fₑₓₜ)
    @named sysₑₓₜ = ODESystem(eqsₑₓₜ, iv, uₑₓₜ, sys.ps)
end
sys_sen = forward_sensitivity_transform(sys,p)

function forward_sensitivity_intial_condition(u₀,p)
    du₀_dp = ModelingToolkit.jacobian(u₀,p)
    u₀ₑₓₜ = vcat(u₀,vec(du₀_dp))
end
x0_sen = forward_sensitivity_intial_condition(x₀,p)

prob_sen = ODEProblem(sys_sen, states(sys_sen).=>x0_sen, tspan, all_parameters .=> rand(length(all_parameters)))
solve(prob_sen,Tsit5())

In constructing the names of the sensitivity states, I had to mess around with the internals quite a bit, and my way is likely not robust. I didn’t keep up with the changes to Symbolics.jl, so there is likely a better way of doing that.

Instead of constructing sensitivity state names with an underscore in them, I would prefer treating them as actual derivative objects. This is where we got stuck previously, as MTK did not have a way of then transforming the ODESystem to an ODEProblem. This seems to still be the case, so finishing the PR will likely take a lot of effort.

1 Like

Honestly, I did not fully grasp the solution you provided, as I am new to the Julia language. I had an alternative idea and wanted to know if it’s a valid approach to perform a forward sensitivity analysis for a subset of parameters.

What if I compute the gradient with respect to each parameter (subset of parameters) for each equation in my system? Then, I thought about adding those equations to the original system, setting zero initial conditions for them, and solving the whole extended system.

using ModelingToolkit, Symbolics, DifferentialEquations

@variables t, x(t), y(t)
@parameters p1, p2, p3, p4

D = Differential(t)

# Original ODE system
eqs = [D(x) ~ p1 * x + p2 * y,
       D(y) ~ p3 * y + p4 * x]

# Differentiate with respect to p1
sensitivity_eqs = [Symbolics.derivative(eq.rhs, p1) for eq in eqs]

# Sensitivity variables
@variables dx_dp1(t), dy_dp1(t)

# Sensitivity equations
sensitivity_system = [D(dx_dp1) ~ sensitivity_eqs[1],
                      D(dy_dp1) ~ sensitivity_eqs[2]]

combined_eqs = vcat(eqs, sensitivity_system)

Sorry for the late reply, I did not get notified of your post.
The sensitivity equations you derive are incorrect. You can find the correct math here:
https://docs.sciml.ai/SciMLSensitivity/stable/sensitivity_math/
Basically, you are not adding the J*S term, and only have the F term.
My piece of code does have both terms in there.