Saving and loading ODEProblems created by ModelingToolkit

Specific Domains Modelling & Simulations
, ,
1

Hello all. I have question about the saving and loading ODEProblems created from Modeling Toolkit (MTK) models. Specifically, when I create and MTK model, build an ODE problem from that, save the ODE problem using JLD2, and then load that ODEProblem, I cannot simulate from it. Here is a MWE.

The snippet of code below creates a simple model (Lotka Volterra), turns it into an ODEProblem, and then saves it with JLD2.

using Pkg; 
Pkg.activate("./PBPK_MTK_demo/");

using DifferentialEquations, ModelingToolkit, OrdinaryDiffEq, JLD2

## Construct MTK model.
LV = function(; name)
	@independent_variables t, [description = "time"]
	Dt = Differential(t)

    pars = @parameters begin
        r = 1.0, [description = "growth rate"]
        K = 2.0, [description = "carrying capacity"]
    end

    vars = @variables begin
        x(t) = 0.1
    end

    eqs = [Dt(x) ~ r*x*(1.0-x/K)]

    ODESystem(eqs, t, vars, pars; name=name)
end

## Instantiate, model and ode problem
@mtkbuild LV1 = LV();

prob = ODEProblem(LV1, [], (0.0, 20.0), []);
jldsave("test_save.jld2", ode_prob = prob, LV1 = LV1);

Next, suppose I kill the Julia instance, start fresh, and import the same packages. If I load the saved ode problem, it will not run.

using DifferentialEquations, ModelingToolkit, OrdinaryDiffEq, JLD2

# Load the JLD2 file
data = jldopen("test_save.jld2", "r") do file
    # Access the stored variables
    ode_prob = file["ode_prob"];
    LV1 = file["LV1"];
    return (ode_prob, LV1)
end;

ode_prob, LV1 = data;

sol = solve(ode_prob, Tsit5(), saveat = 0.1:0.1:20.0);

The solve call generates the following error

KeyError: key (0x5d590ed0, 0x074c081a, 0x531f31a4, 0x516e0718, 0x3a291c4a) not found

However, If I rebuild the ode problem from the MTK model, everything works. Interestingly, the act of constructing a new odeproblem will lead to the old one now solving without issue. It seems something in the MTK model, which is required for the ODEProblem, is lost in the save process.

newprob = ODEProblem(LV1, [], (0.0, 20.0), []);
sol2 = solve(newprob, Tsit5(), saveat = 0.1:0.1:20.0);
sol = solve(ode_prob, Tsit5(), saveat = 0.1:0.1:20.0);

I know this is partly addressed in I/O: Saving and Loading Solution Data · DifferentialEquations.jl . However I am importing OrdinaryDiffEq prior to loading the odeproblem.

It is not technically a big deal to rebuild the ODEProblem every time the MTK model is used in a new script. But that introduces the possibility of human error, simulating a slightly different problem from script to script. It would be ideal to build the ODEProblem once, and then remake it when used in other scripts.

This issue also creates some interactions with Turing.jl when trying to run Bayesian differential equations with a saved ODEProblem (same errors but a bit more complicated workaround). Something seems to be lost in the save → load process. Any thoughts?

Versions used are

Julia 1.11
DifferentialEquations v7.16.1
ModelingToolkit v9.68.1

1 Like
2

IIRC when the ODESystem is generated, ModelingToolkit creates an anonymous function from the expression that gets compiled. That may or may not get serialized correctly. So you may instead want to serialize (eqs, t, vars, pars), and reconstruct the ODEProblem in each Julia session.

3

The runtimegeneratedfunctions need to deserialize as well. This works for Base serialization, so I’d recommend that for now.

Can you open an issue for doing this with other forms of serialization? I’m not sure how to work with JLD2 formats but maybe someone might know how to write a custom read/write hook like that.