How to preserve compilation of observed runtime generated function?

[bradcarman]

I'm afraid my knowledge of RuntimeGeneratedFunction.jl and Julia may not be sufficient to properly ask this question, but let me try my best...

If I run the following code in a fresh Julia session, obviously it will take some time for Julia to compile all the parts needed. If I run it a 2nd time it all runs very quickly, except the sol[] call which always runs slower because of compilation

using ModelingToolkit, OrdinaryDiffEq
include("../examples/rc_model.jl")
sys = structural_simplify(rc_model)
u0 = [
      capacitor.v => 0.0
      capacitor.p.i => 0.0
      resistor.v => 0.0
     ]
prob = ODEProblem(sys, u0, (0, 10.0));
sol = solve(prob, Rodas4());
@time x = sol[resistor.p.i];
@time x = sol[resistor.p.i];

As can be seen, the result of calling sol[resistor.p.i] will always compile the first time if the prob value is newly constructed.

julia> @time x = sol[resistor.p.i];
  0.008358 seconds (9.14 k allocations: 570.337 KiB, 89.54% compilation time)

julia> @time x = sol[resistor.p.i];
  0.000198 seconds (312 allocations: 15.797 KiB)

Why this matters? In my model I have an enormous amount of observed variables thanks to the amazing job of structural_simplify() and this gets my model a very fast run time. But, this efficiency is lost when I extract the solution due to the compilation of the runtime generated observed function. And, if I'm not mistaken, this problem cannot be solved with PackageCompiler.jl because the runtime generated functions are always completely new when an ODEProblem is constructed as far as the compiler understands.

So my question: is there a solution to this problem?

Potential solution 1: If I serialize the prob object, then I can deserialize it and use remake to adjust parameters and tspan. This seems to keep the runtime generated function in the Julia compilation memory which would be stored in the system image by PackageCompiler.jl.

Potential solution 2: Maybe a keyword argument to be supplied to the ODEProblem to store the runtime generated functions to a specified *.jl file. In this way I think the functions would remain more static, even if the ODEProblem is reconstructed. The functions would only need to be recompiled if they change, which I think Julia could keep track of if they are defined in files.

[ChrisRackauckas]

If you need a lot of them, then you can always just use AllObserved() to get the whole array back in one function call.

Solution 1 could work but we would need to get some of the compiler team's time I think 😅, so I'm just going to guess that's not happening soon.

ODEProblemExpr and etc., or expression=Val{:true} in build_function would give back the expressions which you can explicitly eval or save to a file.

[bradcarman]

Thanks Chris, some follow up questions and notes:

1. How do I use AllObserved()? I couldn't find any documentation or examples. My current case I have 300 eliminated variables and I need to grab about 100. So it's hard to say if this will be faster or not. But I'm assuming that the compilation issue will still be there because in the end generate_observed() will still be called (I think).

2. The solution I've landed on is to hack the code generated from build_explicit_observed_function() to build a general function named get_observed(Val(Symbol(obsvar)), u, p, t) which dispatches on the type ::Val{Symbol(obsvar)}

path = ... #generate save path for observable functions...
obs =  ModelingToolkit.observed(sys)
obsvars = [obs[i].lhs for i=1:length(obs)]
for obsvar in obsvars
     obsvarname=replace(string(obsvar), "(t)"=>"") #TODO: need a more robust solution here...
     file = joinpath(path, "$obsvarname.jl")
     obsexp = ModelingToolkit.build_explicit_observed_function(sys, obsvar; expression=true)
     
    # here the function string is modified to provide inputs for multiple dispatch and with function name get_observed()
    obscode = replace(string(obsexp), "function (" => "function get_observed(::Val{Symbol(\"$obsvar\")},")
            
    write(file, obscode)

    include(file)
end

So far this works great! Once the functions are compiled, they remain compiled even if the ODESystem, ODEProblem, etc. are re-generated.

Would love to hear if anyone has a more elegant solution. Not sure if this is already in the plans, but it would be great to serialize the whole ODESystem/ODEProblem to a named function and file with a keyword argument or some other interface. I think the ODEProblemExpr Chris mentions is a good first step, but it wasn't obvious to me how to get from start to finish with that function.

[ChrisRackauckas]

sol[AllObserved()] or sol(t;idxs=AllObserved()) should generate them all in one go.

[ChrisRackauckas]

It's preserved, and new SII handles a lot of this.