SavingCallback when using Ensemble Simulations


I’m able to use SavingCallback with single simulations, for example

using LinearAlgebra
using SparseArrays
using DifferentialEquations

H0 = sprand(ComplexF64, 10, 10, 0.5)
H0 *= H0'
ψ0 = normalize(rand(ComplexF64, 10))

t_l = LinRange(0, 1, 100)

tspan = (t_l[1], t_l[end])

saved_values = SavedValues(Float64, Float64)
save_func = (u, t, integrator) -> norm(u)
cb1 = SavingCallback(save_func, saved_values, saveat = t_l)

dudt!(du,u,p,t) = mul!(du, -1im * H0, u)

prob = ODEProblem(dudt!, ψ0, tspan, callback = cb1)
sol = solve(prob, Tsit5())

100-element Vector{Float64}

But now I need to use the SavingCallback with Ensemble Simulation, for example

function prob_func(prob,i,repeat)
    remake(prob,u0=normalize(rand(ComplexF64, 10)))

prob = ODEProblem(dudt!, ψ0, tspan, callback = cb1)
ensemble_prob = EnsembleProblem(prob, prob_func=prob_func)
sol = solve(ensemble_prob, Tsit5(), EnsembleSerial(), trajectories=10)

But saved_values.saveval remains empty. I noticed that if I move the argument callback = cb1 inside the solve function, saved_values.saveval have the values of the last (I think) simulation.

How can I store the datas of all the simulations in an array of size 100x10 (for this specific case)?

You need to control the deepcopy yourself. I would recommend directly setting the cache in the prob_func.

Ok, so SavingCallback does not support Ensembles yes, and I need to do it by myself, right?

It’s moreso that what you define as saveval is just used as a reference, so you should have one reference per solve if you want it to have different saves. It’s a very general tool of that form.