Allow JumpProblems to be remade with new callback

[sebastiantk]

I have a JumpProblem where I need to save some additional values beyond the state variables. I add a SavingCallback which works great for a single solve.

I then try to create and solve an EnsembleProblem. According to this thread the way to do this is to create a vector of saved_values, one for each trajectory, and provide a prob_func to the EnsembleProblem that remakes the problem for each run with an updated reference to the saved_values vector.

But, I get the error message JumpProblems can currently only be remade with new u0, p, tspan or prob fields. To change other fields create a new JumpProblem. Feel free to open an issue on JumpProcesses to discuss further.

Is it possible to include a new callback to remake with JumpProblems or is there a different way to achieve this?

Thank you!

MWE:

using DifferentialEquations

rate1(u, p, t) = p[1] * u[1] * u[2]
function affect1!(integrator)
    integrator.u[1] -= 1
    integrator.u[2] += 1
end
jump = ConstantRateJump(rate1, affect1!)
u0 = [999, 1, 0]
p = (0.1 / 1000,)
tspan = (0.0, 250.0)
dprob = DiscreteProblem(u0, tspan, p)
jprob = JumpProblem(dprob, Direct(), jump)
saved_vals = SavedValues(Float64, Float64)
cbfun(u, t, integrator) = 2 * u[1]
cb = SavingCallback(cbfun, saved_vals)
jsol = solve(jprob; callback=cb)
saved_vals

ensembleprob = EnsembleProblem(jprob)
ensemblesol = solve(ensembleprob, SSAStepper(); trajectories=10)
ensemblesumm = EnsembleSummary(ensemblesol, quantiles=[0.1, 0.9])

saved_vals_ens = [SavedValues(Float64, Float64) for _ in 1:10]
function prob_func1(prob, i, repeat)
    cb = SavingCallback(cbfun, saved_vals_ens[i])
    remake(prob; callback=cb)
end

ensembleprob2 = EnsembleProblem(jprob, prob_func=prob_func1)
ensemblesol2 = solve(ensembleprob2, SSAStepper(); trajectories=10)

[isaacsas]

This could probably be made to work. As a short-term workaround I'd suggest just calling solve in a loop with a new callback each time, and avoiding EnsembleProblems. Also, if you only care about whatever you are saving via the callback I'd suggest using

jprob = JumpProblem(dprob, Direct(), jump; save_positions = (false, false))

so as not to also save the state during every jump (which can slow the simulations down a lot when running many samples).

[sebastiantk]

Great, thanks, I will try the loop!