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

8

Sorry for recovering this thread, but I am not sure how this approach works when a remake(...) call is to be avoided. As far as I understand, remake always makes a deepcopy of the prob, yet this is something that I would like to avoid as this will create a lot of memory overhead to the point of crashes due to being out of memory.

In essence, this is related to this thread, where I am using TaskLocalValues in order not to deepcopy an ODEProblem (or SDEProblem). But, I cannot find any documentation on how to change the callback in-place, as to avoid using remake.

In essence, the code looks like

tlv_prob = TaskLocalValue{SDEProblem}(() -> deepcopy(prob))

T = Tuple{Vector{Float64}, Vector{Float64}}
saved_values = [SavedValues(Float64, T) for _ in 1:nseeds]

#/ Define prob_func that mutates the (local) SDEProblem
function prob_func(prob, i, nrepeats)
    #~ Get local problem that can be mutated safely
    localprob = tlv_prob[]
    localcallbackset = CallbackSet(savingcallback(saved_values[i]))
    #@TODO Should `remake(...)` be avoided?
    localprob = remake(localprob, callback=localcallbackset)
    return localprob
end

where savingcallback(...) is a function that defines the SavingCallback, which stores some (rolling) mean and variance of a stochastic trajectory. (Note that I use CallbackSet here, as in the full code I have multiple callbacks).
This works fine, but I could’ve just as well avoided the localprob and just use prob instead, as remake initiates a deepcopy of the problem at hand, and I run out of memory sometimes.

I know that the callback of the problem is stored in prob.kwargs.data.callback, but I believe it is stored as an immutable struct and thus it cannot be easily changed. Is there a way, perhaps using replace or something similar, to change the callback in-place, such that it works in conjunction with some SavedValues-arrays? Or can I adjust the saved_values somehow to deal with this?

In other words, is there a setp/setu equivalent for callbacks?

9

It does not. remake(prob, p=p) should just be the same as making a new prob with the new vector swapped in, so just the type wrapper is changed and that should be a cheap operation.

You can use Accessors.jl @set! as a helper.

10

Then why does the doc state:

Calling remake(prob) creates a copy of the existing problem. This new problem has the exact same types as the original one, and the remake call is fully inferred.

Or am I wrong in assuming that a copy and a deepcopy are the same?

And thanks, this looks very useful!

11

Oh, that probably needs to be made more clear. It makes a new object which references the original values and arrays but swaps in the new ones.

12

Ah so in that sense there is essentially no copying unless perhaps for the (symbols of) the state/parameter variables that are changed? That makes using remake perhaps not the problem of the OOM issues that I am experiencing.

13

Yes exactly, it just reuses references to the vectors, so it should be fairly lean. Sounds like we need to make that docstring more clear. But it shouldn’t be the source of your issues here.

Ensembles can deepcopy though, you might want to check safetycopy=false

14

Yep I already put that flag to false. I will investigate more and open a relevant topic if the problem persists. Thank a lot anyways!