ODE parameters estimation using RxInfer

[rzhli]

Parameter estimation using both optimization method and bayesian method are described here, the package the bayesian method used is Turing.jl. I'm wondering weather RxInfer.jl can do ODE parameter estimation, and how to implement it.
Here is the ODE function and data that used:

using OrdinaryDiffEq, StableRNGs

function lotka_volterra!(du, u, p, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = dx = α * x - β * x * y
    du[2] = dy = -δ * y + γ * x * y
end

begin
	lotka_u0 = [1.0, 1.0]
	lotka_p = [1.0, 1.5, 3.0, 1.0]
	lotka_tspan = (0.0, 10.0)
end

lotkaprob = ODEProblem(lotka_volterra!, lotka_u0, lotka_tspan, lotka_p)

lotka_data_sol = solve(lotkaprob, Tsit5(), abstol = 1e-8, reltol = 1e-8, saveat = 0.5)

rng = StableRNG(100)

lotka_data_noisy = abs.(Array(lotka_data_sol) + 0.35randn(rng, size(lotka_data_sol)))

With the ode function and the data(lotka_data_noisy), how to find my parameter lotka_p
Similar example is SIR model as described there, but may be a little different

[bvdmitri]

Hey @rzhli indeed we should add such an example to https://github.com/ReactiveBayes/RxInferExamples.jl
We have some internal code that works but isn't prettified, we will work on it and publish it as a new example and return to you.

[bvdmitri]

Ref: ReactiveBayes/RxInferExamples.jl#8

[rzhli]

Great, I'm looking forward to see it.

[albertpod]

@rzhli
ReactiveBayes/RxInferExamples.jl#9

[bvdmitri]

@rzhli The new example should now be available on the website too https://reactivebayes.github.io/RxInferExamples.jl/dev/examples/Problem%20Specific/ODE%20Parameter%20Estimation/ODE%20Parameter%20Estimation/

[rzhli]

Bravo! Such diverse ways to solve this problem.

One question:
Is noisev necessary for covariance of MvNormalMeanCovariance in each method? Since in real observation we don't know the amplitude of noise.

BTW, can you help me with my attemptation as following:

using OrdinaryDiffEq, StableRNGs
#  Data generation 

function lotka_volterra!(du, u, p, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = dx = α * x - β * x * y
    du[2] = dy = -δ * y + γ * x * y
end

begin
	lotka_u0 = [1.0, 1.0]
	lotka_p = [1.0, 1.5, 3.0, 1.0]
	lotka_tspan = (0.0, 10.0)
end

lotkaprob = ODEProblem(lotka_volterra!, lotka_u0, lotka_tspan, lotka_p)

lotka_data_sol = solve(lotkaprob, Tsit5(), abstol = 1e-8, reltol = 1e-8, saveat = 0.1)

rng = StableRNG(100)

lotka_data_noisy = abs.(Array(lotka_data_sol) + 0.35randn(rng, size(lotka_data_sol)))

function lotka_prediction(u0, p, saveat = 0.5)
	_prob = remake(lotkaprob, u0 = u0, p = p)
	sol = solve(_prob, Tsit5(), saveat, abstol = 1e-8, reltol = 1e-8)
end

using RxInfer

@model function fit_odeparam(obs, m_u, V_u, m_p, V_p, saveat)
    
    σ ~ InverseGamma(2, 3)
    # truncated for avoid using exp function, I'm not sure whether it works in RxInfer
    p ~ truncated(MvNormalMeanCovariance(m_p, V_p);   lower = 0.1, upper = 4.5)

    u0 = MvNormalMeanCovariance(m_u, V_u)

    predicted := lotka_prediction(u0, p, saveat)

    for i in 1:eachindex(predicted)
        obs[:, i] ~ MvNormalMeanCovariance(predicted[:, i], σ^2 * I)  
    end
end

meta = @meta begin
    # Extended Kalman Filter (EKF) type of approximation
    # lotka_prediction -> Linearization() 
    # Or Unscented Kalman Filter (UKF) type of approximation
    lotka_prediction -> Unscented()
end

autoupdates = @autoupdates begin
    m_u, V_u = mean_cov(q(x))
    m_p, V_p = mean_cov(q(p))
end

@initialization function initialize(m_u, V_u, m_p, V_p)
    q(u) = MvNormalMeanCovariance(m_u, V_u)    
    q(p) = MvNormalMeanCovariance(m_p, V_p)     
end

saveat = 0.1

result = infer(
    model = fit_odeparam(saveat = saveat),
    data = (obs = lotka_data_noisy, ),
    initialization = initialize(ones(2), 1e-6diageye(2), zeros(4), 0.1*diageye(4)),
    meta = meta,
    autoupdates = autoupdates,
    showprogress = true,
    free_energy = true
)

Since I'm not familiar with SeeToDee.Rk4 function, I stayed with OrdinaryDiffEq and refered to your examples, but it fails:

ERROR: MixedArguments not supported for rhs_interfaces when node has to be materialized
Stacktrace:

error(s::String) at [error.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

make_node!(::Static.True, node_type::GraphPPL.Atomic, behaviour::GraphPPL.Deterministic, model::GraphPPL.Model{…}, ctx::GraphPPL.Context, options::GraphPPL.NodeCreationOptions{…}, fform::typeof(truncated), lhs_interface::GraphPPL.ProxyLabel{…}, rhs_interfaces::GraphPPL.MixedArguments{…}) at [graph_engine.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

make_node!(atomic::GraphPPL.Atomic, deterministic::GraphPPL.Deterministic, model::GraphPPL.Model{…}, ctx::GraphPPL.Context, options::GraphPPL.NodeCreationOptions{…}, fform::typeof(truncated), lhs_interface::GraphPPL.ProxyLabel{…}, rhs_interfaces::GraphPPL.MixedArguments{…}) at [graph_engine.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

make_node!(::GraphPPL.Atomic, model::GraphPPL.Model{…}, ctx::GraphPPL.Context, options::GraphPPL.NodeCreationOptions{…}, fform::typeof(truncated), lhs_interface::GraphPPL.ProxyLabel{…}, rhs_interfaces::GraphPPL.MixedArguments{…}) at [graph_engine.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

make_node!(model::GraphPPL.Model{…}, ctx::GraphPPL.Context, options::GraphPPL.NodeCreationOptions{…}, fform::typeof(truncated), lhs_interface::GraphPPL.ProxyLabel{…}, rhs_interfaces::GraphPPL.MixedArguments{…}) at [graph_engine.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

macro expansion at [model_macro.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

macro expansion at [ODESystem_in_RxInfer.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

add_terminated_submodel!(__model__::GraphPPL.Model{…}, __context__::GraphPPL.Context, __options__::GraphPPL.NodeCreationOptions{…}, ::typeof(fit_odeparam), __interfaces__::@NamedTuple{…}, ::Static.StaticInt{…}) at [model_macro.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

add_terminated_submodel!(model::GraphPPL.Model{…}, context::GraphPPL.Context, options::GraphPPL.NodeCreationOptions{…}, fform::Function, interfaces::@NamedTuple{…}) at [graph_engine.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

add_terminated_submodel!(model::GraphPPL.Model{…}, context::GraphPPL.Context, fform::Function, interfaces::@NamedTuple{…}) at [graph_engine.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

add_toplevel_model!(model::GraphPPL.Model{…}, context::GraphPPL.Context, fform::Function, interfaces::@NamedTuple{…}) at [graph_engine.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

create_model(callback::RxInfer.var"#24#26"{…}, generator::GraphPPL.ModelGenerator{…}) at [model_generator.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

__infer_create_factor_graph_model(generator::GraphPPL.ModelGenerator{…}, conditioned_on::@NamedTuple{…}) at [model.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

create_model(generator::RxInfer.ConditionedModelGenerator{GraphPPL.ModelGenerator{…}, @NamedTuple{…}}) at [model.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

streaming_inference(; model::GraphPPL.ModelGenerator{…}, data::@NamedTuple{…}, datastream::Nothing, initialization::RxInfer.InitSpecification, autoupdates::RxInfer.AutoUpdateSpecification{…}, constraints::Nothing, meta::GraphPPL.MetaSpecification, options::Nothing, returnvars::Nothing, historyvars::Nothing, keephistory::Nothing, iterations::Nothing, free_energy::Bool, free_energy_diagnostics::Tuple{…}, allow_node_contraction::Bool, autostart::Bool, events::Nothing, addons::Nothing, callbacks::Nothing, postprocess::DefaultPostprocess, uselock::Bool, warn::Bool) at [streaming.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

kwcall(::@NamedTuple{…}, ::typeof(RxInfer.streaming_inference)) at [streaming.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

infer(; model::GraphPPL.ModelGenerator{…}, data::@NamedTuple{…}, datastream::Nothing, autoupdates::RxInfer.AutoUpdateSpecification{…}, initialization::RxInfer.InitSpecification, initmessages::Nothing, initmarginals::Nothing, constraints::Nothing, meta::GraphPPL.MetaSpecification, options::Nothing, returnvars::Nothing, predictvars::Nothing, historyvars::Nothing, keephistory::Nothing, iterations::Nothing, free_energy::Bool, free_energy_diagnostics::Tuple{…}, allow_node_contraction::Bool, showprogress::Bool, catch_exception::Bool, callbacks::Nothing, addons::Nothing, postprocess::DefaultPostprocess, events::Nothing, uselock::Bool, autostart::Bool, warn::Bool) at [inference.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

kwcall(::@NamedTuple{…}, ::typeof(infer)) at [inference.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)

top-level scope at [ODESystem_in_RxInfer.jl](vscode-file://vscode-app/usr/lib/code/out/vs/code/electron-sandbox/workbench/workbench.html)