ModelingToolkit: using it with parallel ensemble simulations

I’m testing out ModelingToolkit to see if I can get a performance improvement. I’ve converted my prob using modelingtoolkitize(prob). I’m wondering if the ModelingToolkit prob is compatible with EnsembleProblem, and what would be the best way to recreate new problems by changing parameters and initial conditions without recalculating the Jacobian each time?

Yes, it’s a normal ODEProblem.

The jacobian is a function of parameters.

Thanks. The interface for passing the variables and parameters to solve is a bit different (it’s a dictionary vs. an array). How would I go about converting an array efficiency into a dict with the variable/parameter keys set up correctly?

EDIT: being more explicit, here’s an example with some questions in comments:

using Distributed
@everywhere using DifferentialEquations
using ModelingToolkit
using Profile
using BenchmarkTools

@everywhere function lorenz!(du,u,p,t)
    du[1] = 10.0*(u[2]-u[1])
    du[2] = u[1]*(28.0-u[3]) - u[2]
    du[3] = u[1]*u[2] - (8/3)*u[3]

u0 = [1.0, 0.0, 0.0]
tspan = (0.0, 100.0)
p = [1, 2.0, 3]
prob = ODEProblem(lorenz!, u0, tspan, p)

sys = modelingtoolkitize(prob)
jac = eval(ModelingToolkit.generate_jacobian(sys)[2])
f = ODEFunction(lorenz!, jac=jac)
prob_jac = ODEProblem(f, u0, tspan, p)

@everywhere function prob_func(prob, i, repeat)
    prob = remake(prob, tspan=(rand(), 100.0), p=rand(3))
    return prob

# version 1 using jacobian and regular interface
ensemble_prob = EnsembleProblem(prob_jac, prob_func = prob_func, safetycopy=true)
@benchmark solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)

# version 2 using modelingtoolkit interface
prob = ODEProblem(sys, u0, tspan, p, jac=true) # this does not work
# Question: how do I get u0 and p into the right data structure given I have used modelingtoolkitize?
ensemble_prob = EnsembleProblem(prob, prob_func = prob_func, safetycopy=true)
# Question: how do I set up prob_func correctly? I'm guessing the current version won't work
@benchmark solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories=100)

My understanding is that although version 1 gives me the Jacobian automatically, version 2 could potentially have more performance enhancements (e.g. index reduction?)

Potentially, but here it won’t do anything. Standard ODEs with non-stiff solvers already written in a way that is non-allocating won’t see a benefit.

u0map = states(sys) .=> u0. We’re going to add a dispatch to make this a little simpler though.

Great, thank you!