Memory usage while solving `JumpProblem` increases with run time, even with `save_positions=(false,false)`

[isaacsas]

@ChrisRackauckas can confirm this pure JumpProcesses version of the first network does not exhibit this issue. So this seems to be coming from the Symbolics codegen or RuntimeGeneratedFunctions.

julia> function lotkavolterradirect(; T = 100.0)
           r1(u,p,t) = growth(p[1], u[1], u[2], p[2]) * u[1]
           r2(u,p,t) = growth(p[1], u[2], u[1], p[2]) * u[2]
           aff1!(int) = int.u[1] += 1
           aff2!(int) = int.u[2] += 1
           p = (.1, 1000)
           u0 = [50, 100]
           tspan = (0.0, T)
           dprob = DiscreteProblem(u0, tspan, p)
           jprob = JumpProblem(dprob, Direct(), ConstantRateJump(r1,aff1!), ConstantRateJump(r2, aff2!); save_positions = (false,false))
           jprob
       end
lotkavolterradirect (generic function with 1 method)

julia> jprob = lotkavolterradirect(T = 10.0)
JumpProblem with problem DiscreteProblem with aggregator Direct
Number of jumps with discrete aggregation: 2
Number of jumps with continuous aggregation: 0
Number of mass action jumps: 0


julia> sol = solve(jprob, SSAStepper())
retcode: Success
Interpolation: Piecewise constant interpolation
t: 2-element Vector{Float64}:
  0.0
 10.0
u: 2-element Vector{Vector{Int64}}:
 [50, 100]
 [102, 216]

julia> @btime solve($jprob, SSAStepper())
  5.858 μs (11 allocations: 1.03 KiB)
retcode: Success
Interpolation: Piecewise constant interpolation
t: 2-element Vector{Float64}:
  0.0
 10.0
u: 2-element Vector{Vector{Int64}}:
 [50, 100]
 [109, 228]

julia> jprob2 = lotkavolterradirect(T = 100.0)
JumpProblem with problem DiscreteProblem with aggregator Direct
Number of jumps with discrete aggregation: 2
Number of jumps with continuous aggregation: 0
Number of mass action jumps: 0


julia> @btime solve($jprob2, SSAStepper())
  30.625 μs (11 allocations: 1.03 KiB)
retcode: Success
Interpolation: Piecewise constant interpolation
t: 2-element Vector{Float64}:
   0.0
 100.0
u: 2-element Vector{Vector{Int64}}:
 [50, 100]
 [402, 598]

Example:

julia> jprob = lotkavolterra(T = 10.0)
JumpProblem with problem DiscreteProblem with aggregator Direct
Number of jumps with discrete aggregation: 2
Number of jumps with continuous aggregation: 0
Number of mass action jumps: 0

julia> aff1 = jprob.discrete_jump_aggregation.affects![1]
RuntimeGeneratedFunction(#=in ModelingToolkit=#, #=using ModelingToolkit=#, :((var"##MTKIntegrator#1026",)->begin
          #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:350 =#
          #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:351 =#
          #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:352 =#
          begin
              ˍ₋out = (var"##MTKIntegrator#1026").u
              ˍ₋arg1 = (var"##MTKIntegrator#1026").u
              ˍ₋arg2 = (var"##MTKIntegrator#1026").p
              t = (var"##MTKIntegrator#1026").t
              begin
                  begin
                      #= /Users/isaacsas/.julia/packages/Symbolics/VIBnK/src/build_function.jl:520 =#
                      #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:399 =# @inbounds begin
                              #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:395 =#
                              ˍ₋out[1] = (+)(1, ˍ₋arg1[1])
                              #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:397 =#
                              nothing
                          end
                  end
              end
          end
      end))

julia> int = (u = [1.0], p = (1.0, 1.0, 1.0), t = 1.0)
(u = [1.0], p = (1.0, 1.0, 1.0), t = 1.0)

julia> @btime aff1($int)
  14.194 ns (1 allocation: 48 bytes)

julia> @btime aff1($int)
  14.194 ns (1 allocation: 48 bytes)

[ChrisRackauckas]

If you take the generated function and just eval it then does it go away?

[isaacsas]

aff1 = jprob.discrete_jump_aggregation.affects![1]
f = eval(aff1)

then gives

julia> @btime aff1($int)
  14.194 ns (1 allocation: 48 bytes)

julia> @btime aff1($int)
  14.194 ns (1 allocation: 48 bytes)

julia> @btime f($int)
  15.573 ns (1 allocation: 48 bytes)

julia> @btime f($int)
  14.194 ns (1 allocation: 48 bytes)

Is that what you meant?

[ChrisRackauckas]

Yeah interesting. Is it returning something? Add a nothing on the bottom to force no return.

[isaacsas]

julia> f = eval(:((var"##MTKIntegrator#1026",)->begin
                        #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:350 =#
                        #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:351 =#
                        #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:352 =#
                        begin
                            ˍ₋out = (var"##MTKIntegrator#1026").u
                            ˍ₋arg1 = (var"##MTKIntegrator#1026").u
                            ˍ₋arg2 = (var"##MTKIntegrator#1026").p
                            t = (var"##MTKIntegrator#1026").t
                            begin
                                begin
                                    #= /Users/isaacsas/.julia/packages/Symbolics/VIBnK/src/build_function.jl:520 =#
                                    #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:399 =# @inbounds begin
                                            #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:395 =#
                                            ˍ₋out[1] = (+)(1, ˍ₋arg1[1])
                                            #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:397 =#
                                            nothing
                                        end
                                end
                            end
                        end; nothing
                    end))
#47 (generic function with 1 method)

julia> @btime f($int)
  14.194 ns (1 allocation: 48 bytes)

julia> @btime f($int)
  14.153 ns (1 allocation: 48 bytes)

julia> function f2(var"##MTKIntegrator#1026")
                              begin
                                   ˍ₋out = (var"##MTKIntegrator#1026").u
                                   ˍ₋arg1 = (var"##MTKIntegrator#1026").u
                                   ˍ₋arg2 = (var"##MTKIntegrator#1026").p
                                   t = (var"##MTKIntegrator#1026").t
                                   begin
                                       begin
                                           #= /Users/isaacsas/.julia/packages/Symbolics/VIBnK/src/build_function.jl:520 =#
                                           #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:399 =# @inbounds begin
                                                   #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:395 =#
                                                   ˍ₋out[1] = (+)(1, ˍ₋arg1[1])
                                                   #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:397 =#
                                                   nothing
                                               end
                                       end
                                   end
                               end
                           end
f2 (generic function with 1 method)

julia> @btime f2($int)
  3.458 ns (0 allocations: 0 bytes)

julia> @btime f2($int)
  3.458 ns (0 allocations: 0 bytes)

Explicitly creating a function with the code does not allocate. Using eval gives allocations with the anonymous function notation it appears.

Note I added a nothing in the anonymous function at the end.

[isaacsas]

Another strange thing is the number of ends that is generated seems unbalanced:

julia> aff1
RuntimeGeneratedFunction(#=in ModelingToolkit=#, #=using ModelingToolkit=#, :((var"##MTKIntegrator#1026",)->begin
          #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:350 =#
          #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:351 =#
          #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:352 =#
          begin
              ˍ₋out = (var"##MTKIntegrator#1026").u
              ˍ₋arg1 = (var"##MTKIntegrator#1026").u
              ˍ₋arg2 = (var"##MTKIntegrator#1026").p
              t = (var"##MTKIntegrator#1026").t
              begin
                  begin
                      #= /Users/isaacsas/.julia/packages/Symbolics/VIBnK/src/build_function.jl:520 =#
                      #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:399 =# @inbounds begin
                              #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:395 =#
                              ˍ₋out[1] = (+)(1, ˍ₋arg1[1])
                              #= /Users/isaacsas/.julia/packages/SymbolicUtils/1JRDc/src/code.jl:397 =#
                              nothing
                          end
                  end
              end
          end
      end))

Notice there are 5 ends but only 4 begins. But dropping one end doesn’t seem to eliminate the issue.

[johannesnauta]

I rewrote my code to not use a Symbolics @reaction_network, similar to your implementation, but I do see the increase in memory allocations still (albeit a bit less)…

Code

(I didn’t want to hide your discussion in my long functions, so here’s a fold:

See code details here

function catalystplasmiddirect(;T = 100.0)    
    function rate(η,X,Y,K)
        return η*(1-(X+Y)/K)
    end

    r1(u,p,t) = rate(p[1],u[1],u[2],p[2])*u[1]
    r2(u,p,t) = rate(p[1],u[2],u[1],p[2])*u[2]
    r3(u,p,t) = p[3]*u[1]
    r4(u,p,t) = p[3]*u[2]
    r5(u,p,t) = p[4]*u[1]*u[2]
    r6(u,p,t) = p[5]*u[2]
    aff1!(integrator) = integrator.u[1] += 1
    aff2!(integrator) = integrator.u[2] += 1
    aff3!(integrator) = integrator.u[1] -= 1
    aff4!(integrator) = integrator.u[2] -= 1
    function aff5!(integrator)
        integrator.u[1] -= 1
        integrator.u[2] += 1
    end
    function aff6!(integrator)
        integrator.u[1] += 1
        integrator.u[2] -= 1
    end
    #    η    K    μ    γ     ρ
    p = (1.0, 1e4, 0.1, 1e-4, 0.01)
    u0 = [1000,10]
    tspan = (0.0, T)

    dprob = DiscreteProblem(u0, tspan, p)
    jprob = JumpProblem(
        dprob, Direct(),
        ConstantRateJump(r1,aff1!), ConstantRateJump(r2,aff2!), ConstantRateJump(r3,aff3!),
        ConstantRateJump(r4,aff4!), ConstantRateJump(r5,aff5!), ConstantRateJump(r6,aff6!);
        save_positions=(false,false)
    )
    return jprob
end

And the benchmarks:

julia> long_prob = catalystplasmiddirect(T=1e5)
JumpProblem with problem DiscreteProblem with aggregator Direct
Number of jumps with discrete aggregation: 6
Number of jumps with continuous aggregation: 0
Number of mass action jumps: 0


julia> long_sol = solve(long_prob, stepper)
retcode: Success
Interpolation: Piecewise constant interpolation
t: 2-element Vector{Float64}:
      0.0
 100000.0
u: 2-element Vector{Vector{Int64}}:
 [1000, 10]
 [98, 8850]

julia> @btime solve($long_prob, $stepper)
  11.082 s (195816952 allocations: 2.92 GiB)
retcode: Success
Interpolation: Piecewise constant interpolation
t: 2-element Vector{Float64}:
      0.0
 100000.0
u: 2-element Vector{Vector{Int64}}:
 [1000, 10]
 [96, 8951]

julia> short_prob = catalystplasmiddirect(T=1e2)
JumpProblem with problem DiscreteProblem with aggregator Direct
Number of jumps with discrete aggregation: 6
Number of jumps with continuous aggregation: 0
Number of mass action jumps: 0

julia> short_sol = solve(short_prob, stepper)
retcode: Success
Interpolation: Piecewise constant interpolation
t: 2-element Vector{Float64}:
   0.0
 100.0
u: 2-element Vector{Vector{Int64}}:
 [1000, 10]
 [92, 8892]

julia> @btime solve($short_prob, $stepper)
  10.677 ms (204394 allocations: 3.12 MiB)
retcode: Success
Interpolation: Piecewise constant interpolation
t: 2-element Vector{Float64}:
   0.0
 100.0
u: 2-element Vector{Vector{Int64}}:
 [1000, 10]
 [118, 8933]

[isaacsas]

@johannesnauta

short_prob = TestJumpProcesses.catalystplasmiddirect(T=1e2)

OK, but if you switch to DirectFW, which uses FunctionWrappers, the memory allocations stay the same for me (make sure you call the same function to generate your problem for short and long times).

So the issue in this specific case is from what I said earlier with Direct calling the affects array and creating a union type I think. However, there is still an issue with Symbolics generated code, which seems to still exhibit allocations even when using DirectFW.

[johannesnauta]

Indeed, when I change Direct() to DirectFW() the large memory allocations disappear. Is this documented somewhere?

And there is indeed still the issue with Symbolics generated code. I can confirm the increased memory usage even when using DirectFW(). Thanks anyways for highlighting the other approach(es) to implementing reactions.

[isaacsas]

Any SSA method but Direct will use FunctionWrappers internally and so shouldn’t have this issue (when directly constructing the jumps – the issue with codegen from Symbolics is different).

For your problem, if you build a dependency graph as explained in the JumpProcesses docs I’d guess one of SortingDirect, RSSA or RSSACR is going to be fastest (depending on how big your full model is).

[isaacsas]

I should add, if you find allocation growth on these examples with other methods (when building the problem directly in JumpProcesses) please do let me know.

[johannesnauta]

I originally used the RSSA() method (but with the Symbolics codegen) and noticed the growth. I will try to figure out how to define a MassActionJump with a variable rate (should I use a VariableRateJump, or does that one not work with the RSSA() method?).

(CODE) With a fixed rate it's not hard to create the stoichiometry.

function catalystplasmidrssa(;T = 100.0)
    rateidxs = [1,1,2,2,3,4]
    reactant_stoich = [
        [1 => 1],  # 1*F  (reproduction)
        [2 => 1],  # 1*P  (reproduction)
        [1 => 1],  # 1*F  (death)
        [2 => 1],  # 1*P  (death)
        [2 => 1],  # 1*P  (recovery, segregation error)
        [1 => 1, 2=> 1]  # 1*F 1*P  (infection, conjugation)
    ]
    net_stoich = [
        [1 => 1],  # F --> 2F
        [2 => 1],  # P --> 2P
        [1 => -1], # F --> 0
        [2 => -1], # P --> 0
        [1 => 1, 2 => -1],  # P --> F
        [1 => -1, 2 => 1]   # F+P --> 2P
    ]
    
    u0 = [1000,10]
    tspan = (0.0, T)
    K = 1e4
    #    η    μ    ρ    γ
    p = (1.0, 0.1, 0.1, 1e-4)

    dprob = DiscreteProblem(u0, tspan, p)
    mass_action_jump = MassActionJump(reactant_stoich, net_stoich; param_idxs=rateidxs)
    jump_prob = JumpProblem(dprob, RSSA(), mass_action_jump; save_positions=(false,false))
    return jump_prob
end

As the carrying capacity term is not included this leads to infinite growth (i.e. depending on \gamma we have either F\rightarrow \infty or P\rightarrow \infty, so do not choose large (\gtrsim 10^2) timespans or your REPL will hang due to the integers becoming very large very quick.

But, as mentioned in the links above, I see nothing in the tutorial about variable rates outside of the example that I provided earlier with DirectFW(). So for larger systems I will still need to figure it out. Once I do I will make a nice example to add to the documentation perhaps.

Also, if this needs to be made into an issue (either for Catalyst.jl or for JumpProcesses.jl?) for better tracking, let me know as well.

[isaacsas]

@johannesnauta have you gone through the JumpProcesses tutorials? You want to use a mix of MassActionJumps and ConstantRateJumps for your problem. The latter let you have arbitrary rate functions as long as they are not explicitly time-dependent so can be used to model a carrying capacity. For a pure jump problem you would only use VariableRateJumps if your rate expressions explicitly include t. See for example:

https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/

[johannesnauta]

Yes I went through the JumpProcesses tutorial (more specifically this part, which is why I used the ConstantRateJumps as illustrated above in the code details). However, as you mentioned, that code used Direct() instead of a method that uses FunctionWrappers.
While it worked fine when using DirectFW(), it does not work with other methods, such as RSSA() as:

ERROR: To use the RSSA algorithm a map from variables to dependent jumps must be supplied.

So I was under the impression I needed to implement the stoichiometry directly, as explained in the MassActionJump documentation/example.

Edit: also there is a note on this error here: Jump types and JumpProblem · JumpProcesses.jl, but I am unsure how to implement it from that documentation alone.

[isaacsas]

Can you try Direct with JumpProcesses 9.6.1? I put in what seems to be a fix for the memory issue (when directly constructing jumps, it doesn’t fix the issue with the codegen from Symbolics unfortunately).

[isaacsas]

Regarding dependency graphs, for a reaction-reaction dependency graph you need a mapping from the index representing the ith reaction to a vector of the indices of all reactions that depend on the ith reaction. Basically this is a mapping that says when reaction i occurs, which reactions have rate functions that depend on a species that is modified by reaction i.

Reaction indices are ordered internally by the mass action jumps (in the order they are passed to MassActionJump) and then by ConstantRateJumps in the order they are passed, see this example for how the ordering should work:

“FAQ · JumpProcesses.jl”

[isaacsas]

@johannesnauta I’ve also checked that with the new Direct in JumpProcesses 9.6.1 your Catalyst model from catalyst_plasmid does not appear to show the memory growth anymore (but there is still growth for other solvers).

[johannesnauta]

Sorry for the late reply. I was in the middle of moving and things did not go as planned.

But, with v.9.6.1 of JumpProcesses the catalyst_plasmid function indeed does not show the memory growth anymore. In fact, the overall memory usage has decreased drastically in general. I will further explore the functions that define the MassActionJumps and ConstantRateJumps directly, but checking here I believe that Direct() will do for now – until I scale my problem to larger networks. Even still, I do not think I will have many updates per jump, so perhaps Direct() will remain the fastest method.

Thanks a lot for the effort and big apology for the long delay.

[isaacsas]

I have a PR that appears to fix the issues for all the non-direct methods for models coming from Catalyst too:

function wrapper type instability fix

(It also gives a 10-25% speedup for the other methods courtesy of the fix it appears.)

I probably won’t get it done this week, but hopefully next. In any case, once that gets merged you should be able to use Catalyst-generated models with any of the constant rate aggregators without issue. I’ll let you know once a new release with those fixes is out.

[isaacsas]

Ok, this should now be fixed for all methods as of version 9.6.3.