Simulating a large number of independent jumps

[pooyaravari]

We want to simulate a large number of independent jumps in a SDE problem. Code is below. When N = 300, it's reasonable, however for N = 1000 (for some cases we might need N = 10000) it's stalling. As discussed in SciML/DiffEqDocs.jl#258 (comment)
SSAStepper() is no use here, but even if wrongly someone use it, with N = 1000 it's stalling.

This type of model we are using is for continuous-time models of an economy with heterogeneous firms.

Here is a sample code:

using DifferentialEquations ,BenchmarkTools,  StochasticDiffEq, DiffEqJump, Plots

struct AffectIndex{F1, F2}
  affect_index!::F1
  index::F2
end
function (p::AffectIndex)(integrator)
    return p.affect_index!(integrator, p.index)
end

####
function μ_SDE(du,u,p,t)
  du .= p.μ
end

function σ_SDE(du,u,p,t)
  du .= p.σ
end

function test_SRIW1()
    p = (μ = 0.01, σ = 0.1, N = 500) # if all constant
    T = 10.0  # maximum time length
    x_iv = rand(p.N)  # just draws from the inital condition

    prob = SDEProblem(μ_SDE, σ_SDE, x_iv ,(0.0, T), p)
    rate(u,p,t) = 0.2
    affect_index!(integrator, index) = (integrator.u[index] = max(integrator.u[index], integrator.u[rand(1:integrator.p.N)]) )
    jumps = [ConstantRateJump(rate,AffectIndex(affect_index!, i)) for i in 1:p.N]
    jump_prob = JumpProblem(prob,Direct(),JumpSet((), jumps, nothing, nothing))
    @btime solve($jump_prob, $SRIW1());
end

[pooyaravari]

Also if want to check with closures method, here is the code:

function test_closures()
    p = (μ = 0.01, σ = 0.1, N = 300) # if all constant
    T = 10.0  # maximum time length
    x_iv = rand(p.N)  # just draws from the inital condition

    prob = SDEProblem(μ_SDE, σ_SDE, x_iv ,(0.0, T), p)
    rate(u,p,t) = 0.2
    affect_index!(integrator, index) = (integrator.u[index] = max(integrator.u[index], integrator.u[rand(1:integrator.p.N)]) )
    jumps = [ConstantRateJump(rate,AffectIndex(affect_index!, i)) for i in 1:p.N]
    jump_prob = JumpProblem(prob,Direct(),jumps...)
    @btime solve($jump_prob);
end
test_closures()

[isaacsas]

Thanks, I’ll play around with this when I’m back in Boston next week.

[jlperla]

Thanks. We will investigate other approaches for when the jumps are all of identical frequency (as discussed before), so the main purpose for this would be to have a framework to modify for when the jumps have different rates

[jlperla]

But also, we just aren't sure what to consider swapping out for the Direct() and SRIW1() in general. Our understanding is that changjng Direct to DirectFW in the closure example is the best choice, but otherwise we don't know permutations of either.

[ChrisRackauckas]

Direct() and SRIW1 are completely different things and cannot be swapped. SSAStepper is a DiscreteProblem-only fast stepper which is an alternative to SRIW1 if there's no differential equation. Direct is an aggregator, which would need to be swapped with a different jump aggregator.

[jlperla]

OK. So it sounds like the following is true: Direct -> DirectFW if we are wrapping a whole bunch of things of different types, otherwise it shouldn't matter.
And SRIW1 is the only alternative with this setup (since we aren't a pure jump process).

[ChrisRackauckas]

SOSRI is a bit better.

[isaacsas]

@jlperla Just an update. I've been playing around with this, and I am seeing that even with save_positions=(false,false) and saving at a small number of time points the performance is pretty slow (N=300). Unfortunately the current brokenness of Cairo.jl means that I haven't been able to actually visualize any profiling data to see what is going on.

Did you or @PooyaFa ever profile the code and see where it is actually spending its time? (i.e. is the major amount of time actually being spent in Direct, or is it in another part of the DiffEq library code?)

[jlperla]

Thanks! I think based on what we discussed with Chris, tau leaping seems like the only way to really jack up the number of points. However, even without the jump size issues, the current regular jump code only does pure jump processes. Alas, we need jump diffusions.

[isaacsas]

I don't think tau-leaping is needed, and it can potentially introduce other issues. One should certainly be able to simulate thousands of jump processes coupled to diffusion processes without needing tau-leaping. The problem is determining where the bottleneck in the current code base is. It could be in Direct, but it could also be in how mixed jump-diffusion systems are being handled by the timestepper. Profiling, if Cairo.jl ever gets fixed, would help in determining where the bottleneck is occurring.

[isaacsas]

If I've understood your problem / code correctly, there are no "dependencies" between the jumps. That is, when a given jump occurs it does not change u in such a way that we would need to recalculate the rate for other jumps (since none of the jump rates actually depend on u). In that case NRM should be faster and work well.

Could you try this code out, which uses NRM, and see if it is actually giving results that look accurate to you?

function test_SRIW1()
           p = (μ = 0.01, σ = 0.1, N = 500) # if all constant
           T = 10.0  # maximum time length
           x_iv = rand(p.N)  # just draws from the inital condition

           prob = SDEProblem(μ_SDE, σ_SDE, x_iv ,(0.0, T), p)
           rate(u,p,t) = 0.2
           affect_index!(integrator, index) = (integrator.u[index] = max(integrator.u[index], integrator.u[rand(1:integrator.p.N)]) )
           jumps = [ConstantRateJump(rate,AffectIndex(affect_index!, i)) for i in 1:p.N]
           jump_prob = JumpProblem(prob,NRM(),JumpSet((), jumps, nothing, nothing),save_positions=(false,false),dep_graph=[Int[] for i=1:p.N])
           sol = solve(jump_prob, SRIW1(), saveat=T);
           #@btime begin solve($jump_prob, $SRIW1(), saveat=$T); end;
           return sol;
       end

[isaacsas]

Note that it currently only saves at the final time point, but you could add more save points using saveat.

[isaacsas]

Also, did you ever actually try DirectFW with save_positions=(false,false) and specifying some saveat? I'm finding just this is much faster on my computer. Direct is 12.241 s, DirectFW is 91.373 ms and NRM is 74.197 ms.

[jlperla]

I am going to close this out for now, as I think that the combination of not saving everything and using DirectFW (or NRM) closes this particular issue.