Parallel Ensemble Simulations

Parallel Ensemble Simulations#

Docs: https://diffeq.sciml.ai/stable/features/ensemble/

Solving an ODE With Different Initial Conditions#

Solving \(\dot{u} = 1.01u\) with \(u(0)=0.5\) and \(t \in [0, 1]\)

using DifferentialEquations
using Plots

[ Info: Precompiling IJuliaExt [2f4121a4-3b3a-5ce6-9c5e-1f2673ce168a]

Linear ODE which starts at 0.5 and solves from t=0.0 to t=1.0

prob = ODEProblem((u, p, t) -> 1.01u, 0.5, (0.0, 1.0))

ODEProblem with uType Float64 and tType Float64. In-place: false
timespan: (0.0, 1.0)
u0: 0.5

Define a new problem for each trajectory using remake() The initial conditions (u0) are changed in this example

function prob_func(prob, i, repeat)
    remake(prob, u0=rand() * prob.u0)
end

prob_func (generic function with 1 method)

You could also obtain the necessary data from the outside of prob_func()

initial_conditions = range(0, stop=1, length=100)

function prob_func(prob, i, repeat)
  remake(prob, u0=initial_conditions[i])
end

Define an ennsemble problem

ensemble_prob = EnsembleProblem(prob; prob_func=prob_func)

EnsembleProblem with problem ODEProblem

Ensemble simulations use multithreading by default

sim = solve(ensemble_prob, trajectories=100)

EnsembleSolution Solution of length 100 with uType:
ODESolution{Float64, 1, Vector{Float64}, Nothing, Nothing, Vector{Float64}, Vector{Vector{Float64}}, ODEProblem{Float64, Tuple{Float64, Float64}, false, SciMLBase.NullParameters, ODEFunction{false, SciMLBase.AutoSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{0, Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, Vern7{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, Rosenbrock23{0, false, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, Rodas5P{0, false, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, FBDF{5, 0, false, Nothing, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, Nothing, Nothing, typeof(OrdinaryDiffEq.trivial_limiter!)}, FBDF{5, 0, true, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, Nothing, Nothing, typeof(OrdinaryDiffEq.trivial_limiter!)}}, OrdinaryDiffEq.AutoSwitchCache{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, Vern7{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}}, Tuple{Rosenbrock23{0, false, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, Rodas5P{0, false, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, FBDF{5, 0, false, Nothing, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, Nothing, Nothing, typeof(OrdinaryDiffEq.trivial_limiter!)}, FBDF{5, 0, true, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, Nothing, Nothing, typeof(OrdinaryDiffEq.trivial_limiter!)}}, Rational{Int64}, Int64}}, OrdinaryDiffEq.InterpolationData{ODEFunction{false, SciMLBase.AutoSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{Float64}, Vector{Float64}, Vector{Vector{Float64}}, Vector{Int64}, OrdinaryDiffEq.DefaultCache{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Vern7ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{_A, SciMLBase.TimeDerivativeWrapper{false, ODEFunction{false, SciMLBase.AutoSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, SciMLBase.NullParameters}, SciMLBase.UDerivativeWrapper{false, ODEFunction{false, SciMLBase.AutoSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, SciMLBase.NullParameters}, Float64, OrdinaryDiffEq.StaticWOperator{true, Float64}, Nothing, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}}} where _A, OrdinaryDiffEq.Rosenbrock5ConstantCache{SciMLBase.TimeDerivativeWrapper{false, ODEFunction{false, SciMLBase.AutoSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, SciMLBase.NullParameters}, SciMLBase.UDerivativeWrapper{false, ODEFunction{false, SciMLBase.AutoSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, SciMLBase.NullParameters}, Tab, Float64, OrdinaryDiffEq.StaticWOperator{true, Float64}, Nothing} where Tab<:OrdinaryDiffEq.Rodas5Tableau, OrdinaryDiffEq.FBDFConstantCache{5, var"#5493#N", Vector{Float64}, Float64, Matrix{Float64}, Matrix{Float64}, StaticArraysCore.SMatrix{5, 6, Rational{Int64}, 30}, _A, Vector{Float64}, Vector{Float64}} where {var"#5493#N"<:(OrdinaryDiffEq.NLSolver{NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, false, Float64, _A, Nothing, _B, OrdinaryDiffEq.NLNewtonConstantCache{tType, tType2, J, W, ufType}} where {_A, _B, tType, tType2, J, W, ufType}), _A}, OrdinaryDiffEq.FBDFConstantCache{5, var"#5493#N", Vector{Float64}, Float64, Matrix{Float64}, Matrix{Float64}, StaticArraysCore.SMatrix{5, 6, Rational{Int64}, 30}, _A, Vector{Float64}, Vector{Float64}} where {var"#5493#N"<:(OrdinaryDiffEq.NLSolver{NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, false, Float64, _A, Nothing, _B, OrdinaryDiffEq.NLNewtonConstantCache{tType, tType2, J, W, ufType}} where {_A, _B, tType, tType2, J, W, ufType}), _A}, Tuple{Float64, Float64, DataType, DataType, DataType, Float64, Float64, ODEFunction{false, SciMLBase.AutoSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64, Float64, Float64, SciMLBase.NullParameters, Bool, Val{false}}, OrdinaryDiffEq.AutoSwitchCache{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, Vern7{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}}, Tuple{Rosenbrock23{0, false, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, Rodas5P{0, false, Nothing, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!)}, FBDF{5, 0, false, Nothing, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, Nothing, Nothing, typeof(OrdinaryDiffEq.trivial_limiter!)}, FBDF{5, 0, true, KrylovJL{typeof(Krylov.gmres!), Int64, Tuple{}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, typeof(OrdinaryDiffEq.DEFAULT_PRECS), Val{:forward}, true, nothing, Nothing, Nothing, typeof(OrdinaryDiffEq.trivial_limiter!)}}, Rational{Int64}, Int64}}, Nothing}, SciMLBase.DEStats, Vector{Int64}, Nothing, Nothing}

Each element in the result is an ODE solution

sim[1]

retcode: Success
Interpolation: 3rd order Hermite
t: 5-element Vector{Float64}:
0
09965328584043498
3457383158593924
677757532821908
0
u: 5-element Vector{Float64}:
39400376596104825
4357245537489117
5586684392089628
7812502708659095
0817770872731949

You can plot all the results at once

plot(sim, linealpha=0.4)

Solving an SDE with Different Parameters#

function lotka_volterra!(du, u, p, t)
    du[1] = p[1] * u[1] - p[2] * u[1] * u[2]
    du[2] = -3 * u[2] + u[1] * u[2]
end

function g!(du, u, p, t)
    du[1] = p[3] * u[1]
    du[2] = p[4] * u[2]
end

p = [1.5, 1.0, 0.1, 0.1]
prob = SDEProblem(lotka_volterra!, g!, [1.0, 1.0], (0.0, 10.0), p)

SDEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 10.0)
u0: 2-element Vector{Float64}:
 1.0
 1.0

function prob_func(prob, i, repeat)
    x = 0.3 * rand(2)
    remake(prob, p=[p[1:2]; x])
end

ensemble_prob = EnsembleProblem(prob, prob_func=prob_func)
sim = solve(ensemble_prob, SRIW1(), trajectories=10)

fig = plot(sim, linealpha=0.6, color=:blue, idxs=(0, 1))
plot!(fig, sim, linealpha=0.6, color=:red, idxs=(0, 2))

Show the distribution of the ensemble solutions

summ = EnsembleSummary(sim, 0:0.1:10)
plot(summ, fillalpha=0.5)

Ensemble simulations of Modelingtoolkit (MTK) models#

Radioactive decay example

using ModelingToolkit
using DifferentialEquations

@variables t c(t) = 1.0
@parameters λ = 1.0
D = Differential(t)
@mtkbuild sys = ODESystem([D(c) ~ -λ * c], t)
prob = ODEProblem(sys, [], (0.0, 2.0), [])

ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 2.0)
u0: 1-element Vector{Float64}:
 1.0

Use the symboilic interface to change the parameter(s) By prob.ps[symbol] or setp() function factory

function changemtkparam(prob, i, repeat)
    # Make a new copy of the parameter vector
    # Ensure the changed will not affect the original ODE problem
    newprob = remake(prob, p=copy(prob.p))
    newprob.ps[λ] = i * 0.5
    newprob
end

ensemble_prob = EnsembleProblem(prob, prob_func=changemtkparam)
sim = solve(ensemble_prob, trajectories=10)
plot(sim, fillalpha=0.5)

This notebook was generated using Literate.jl.

Parallel Ensemble Simulations

Contents

Parallel Ensemble Simulations#

Solving an ODE With Different Initial Conditions#

Solving an SDE with Different Parameters#

Ensemble simulations of Modelingtoolkit (MTK) models#