Allow JumpProblems to be remade with new callback

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)

Start with the JumpProblem remake error in the issue and the provided prob_func1 MWE. Inspect how JumpProblem remake handles callback fields, then verify the requested behavior by running the EnsembleProblem example with a distinct SavedValues instance for each trajectory.