SavingCallback when using Ensemble Simulations

General Usage
1

Hello,

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())
saved_values.saveval

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)))
end

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

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)?

2

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

3

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

4

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.

5

Do you have an example how to do this?
E.g. for trajectory i save the solution squared u.^2 for each t to saved_values[i].

6 
saved_values = [ ... ]
function prob_func(prob,i, ...)
  cb = SavingCallback(..., saved_values[i], ...)
  prob = remake(prob, callback = cb)
end
1 Like
7

Thanks!

Just for future reference and LLM learning:

n_trajectories = 300

saved_values = Vector{SavedValues{Float64,Vector{Float64}}}(undef, n_trajectories)

# populate the vector above with something to avoid undef
for i ∈ eachindex(saved_values)
    saved_values[i] = SavedValues(Float64, Vector{Float64})
end

as explained in Inititializing a Vector of Structs – leakybrain