Increase perfomance for Gillespie Algorithm

Luis · December 1, 2021, 12:40pm

I am relatively new to Julia and trying to implement a version of Gillespies Algorithm that can be used in population genetics.

I tried really hard to make it as fast and light weighted as possible, incorporating most of the performance tips from the Julia documentation. But despite my effort I seemed to achive the exact opposite: A very slow and memory gulping algorithm.

Unfortunately I was not even able to isolate the problem. Therefore I am very happy if anyone is interested in investing the time to look over my code and point out the memory and speed killer I implemented. Or if anyone can hand me some tools or literature that enables me to isolate the problem myselfe.

You can find all the code including some working (but slow) examples in the following git repository:

https://github.com/roccminton/Gillespie_Algorithm

Many thanks in advance for any help!

josuagrw · December 1, 2021, 12:41pm

Unfortunately, the link leads to a 404 page

Luis · December 1, 2021, 12:47pm

Posting a link to a private repsoitory is not very clever. Sorry for that. I changed the visibility to public. Does the link work now.

josuagrw · December 1, 2021, 3:04pm

Yes it works now

Oscar_Smith · December 1, 2021, 3:34pm

Can you profile this and figure out what is taking time? Debugging performance of MWEs is much easier than a full repo.

goerch · December 1, 2021, 3:46pm

The only example working for me on 1.8.0 seems to be Examples\RandomMating.jl. Is that what you are after?

Luis · December 1, 2021, 4:33pm

I fixed the bugs making the other examples not work, but indeed the Examples\RandomMating.jl is the most interesting example for me.

Luis · December 1, 2021, 4:37pm

Try to profile for example the function DiploidModel.rungillespie(t,x0,model_parameter) in the file Examples\RandomMating.jl (in line 28). I did as well but was unable to get helpful information out of the print output.

goerch · December 1, 2021, 4:49pm

After disabling ProgressMeter and Plots and putting the test into a function like so

function test()

        K = 1000

        b = 1.0
        d = 0.9
        c = (b-d)/(2K)

        dni = 0.1
        N = 10

        model_parameter = (
                birth = b,
                death = d,
                competition = c,
                μ = dni,
                Nloci = N
                )

        t = 0:100
        x0 = DiploidModel.generatehealthypopulation(K,N)

        #execute the simulation
        history = DiploidModel.rungillespie(t,x0,model_parameter)
        #plot simulation
        #PlotFromDicts.plotmutationloadandprevalence(history)
end

JET reports for @report_opt test()

═════ 3 possible errors found ═════
┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\Examples\RandomMating.jl:30 Main.DiploidModel.rungillespie(t, x0, model_parameter)
│┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:35 Main.DiploidModel.setupparameter(model_parameter, n₀)
││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:51 Main.DiploidModel.npropagable(n0)
│││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:95 Main.DiploidModel.sum(Base.Generator(#3, ps))
││││┌ @ reduce.jl:549 Base.#sum#262(Base.pairs(Core.NamedTuple()), #self#, a)
│││││┌ @ reduce.jl:549 Base.sum(Base.identity, a)
││││││┌ @ reduce.jl:520 Base.#sum#261(Base.pairs(Core.NamedTuple()), #self#, f, a)
│││││││┌ @ reduce.jl:520 Base.mapreduce(f, Base.add_sum, a)
││││││││┌ @ reduce.jl:294 Base.#mapreduce#258(Base.pairs(Core.NamedTuple()), #self#, f, op, itr)
│││││││││┌ @ reduce.jl:294 Base.mapfoldl(f, op, itr)
││││││││││┌ @ reduce.jl:162 Base.#mapfoldl#254(Base._InitialValue(), #self#, f, op, itr)
│││││││││││┌ @ reduce.jl:162 Base.mapfoldl_impl(f, op, init, itr)
││││││││││││┌ @ reduce.jl:44 Base.foldl_impl(op′, nt, itr′)
│││││││││││││┌ @ reduce.jl:49 Base.reduce_empty_iter(op, itr)
││││││││││││││┌ @ reduce.jl:370 Base.reduce_empty_iter(op, itr, Base.IteratorEltype(itr))
│││││││││││││││┌ @ reduce.jl:371 Base.reduce_empty(op, Base.eltype(itr))
││││││││││││││││┌ @ reduce.jl:348 Base.mapreduce_empty(#3, $(QuoteNode(Base.BottomRF{typeof(Base.add_sum)}(Base.add_sum))), _)
│││││││││││││││││ runtime dispatch detected: Base.mapreduce_empty(#3, $(QuoteNode(Base.BottomRF{typeof(Base.add_sum)}(Base.add_sum)))::Base.BottomRF{typeof(Base.add_sum)}, _::Type{Pair{Main.DiploidModel.RandomIndividual{Int64}, Int64}})
││││││││││││││││└─────────────────
│┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:35 Main.DiploidModel.Gillespie.run_gillespie!(time, n₀, Main.DiploidModel.setupparameter(model_parameter, n₀), Main.DiploidModel.execute!, Main.DiploidModel.rates!, initrates, population_history)
││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:17 Main.DiploidModel.Gillespie.mainiteration!(population_history, initrates, n₀, Main.DiploidModel.Gillespie.convert(Main.DiploidModel.Gillespie.Float64, Base.getindex(time, 1)), time, par, execute!, rates!)
│││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:64 ct = Main.DiploidModel.Gillespie.onestep!(n0, rates, ct, step, par, ex!, r!)
││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:147 Main.DiploidModel.Gillespie.iszero(i)
│││││┌ @ number.jl:42 Base.zero(x)
││││││ runtime dispatch detected: Base.zero(x::Nothing)
│││││└────────────────
││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:151 ex!(i, x_0, par)
│││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:194 Main.DiploidModel.death!(ps, par)
││││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:187 Base.setindex!(ps, Main.DiploidModel.-(Base.getindex(ps, Main.DiploidModel.choosefey(ps, Base.getindex(Base.getproperty(par, :popsize), 1))), Main.DiploidModel.one(Main.DiploidModel.valtype(ps))), Main.DiploidModel.choosefey(ps, Base.getindex(Base.getproperty(par, :popsize), 1)))
│││││││┌ @ dict.jl:374 Base.convert(_, key0)
││││││││ runtime dispatch detected: Base.convert(_::Type{Main.DiploidModel.RandomIndividual{Int64}}, key0::Nothing)
│││││││└───────────────

Does this help?

ChrisRackauckas · December 1, 2021, 4:55pm

How does it compare to the DifferentialEquations.jl/Catalyst.jl SSAStepper?

https://benchmarks.sciml.ai/html/Jumps/NegFeedback_GeneExpr.html

https://catalyst.sciml.ai/dev/tutorials/using_catalyst/

It would be good to at least get a baseline.

ChrisRackauckas · December 1, 2021, 4:56pm

Or use the same baseline as these benchmarks:

https://github.com/sdwfrost/Gillespie.jl#benchmarks

goerch · December 1, 2021, 5:06pm

I’m spotting at least one problem here:

function chooseevent(rates,total_rate)
    #make it a uniform random variable in (0,total_rate)
    rndm = rand(Uniform(0.0,total_rate))
    #choose the rate at random
    @inbounds for (index, rate) in enumerate(rates)
        rndm -= rate
        rndm ≤ 0.0 && return index
    end
end

chooseevent could return nothing

isaacsas · December 1, 2021, 5:20pm

@Luis can you point to any mathematical write ups of one of your models showing how they differ via the trait space from traditional Gillespie models? That might help in pointing you towards possible algorithmic improvements and/or seeing how to add your models into existing packages.

How many rates do you typically sum up each step of the algorithm to get the total rate? If it is many there are lots of possible optimized Gillespie methods one can use, typically designed around only updating individual rates that would have actually changed after the occurrence of a given jump.

adolgert · December 2, 2021, 12:10am

Maybe I can give two pointers here, to papers that discuss what @isaacsas is describing.

Li H, Petzold L (2006) Logarithmic direct method for discrete stochastic simulation of chemically reacting systems. Technical report. Available: http://www.engineering.ucsb.edu/cse/Files/ldm0513.pdf. Accessed on 2012 Apr 24. - This is the standard for optimized Gillespie methods. It uses a prefix sum tree (like a Fenwick tree) to keep track of the sum of the propensities, while allowing you to modify the values of those propensities.
https://www.researchgate.net/publication/8379343_Efficient_formulation_of_the_stochastic_simulation_algorithm_for_chemically_reacting_systems - This is an overview of techniques, showing the scope of the kinds of solutions.

If you want to do a Gillespie algorithm that has an infinite number of possible reactions, with only a finite number active at any one time, then that’s also totally possible, but it just tends to go by other names than “Gillespie.” That’s also in the second reference, above. It entails keeping a mutable heap of the next reactions to fire and the times at which they would fire.

Luis · December 2, 2021, 8:56am

Many thanks for your effort. That is much more readable then the output I usually get!
The first error comes during the setup of the model and gets called only once, thus does not contribute too much to the problem.
Error two and three appear due to the same problem @goerch pointed out. Eventhough the function will never actually return nothing, because the handed in total_rate is the sum of rates I added another return statement at the end of the loop.

Luis · December 2, 2021, 9:16am

I am analyzing adaptive dynamics models as described in Stochastic models for adaptive dynamics: Scaling limits and diversity. There usually the trait space is of infinitely large (e.g. [-1,1] as in the DieckmannDoeblin example) or very very big (e.g. {0,1,2}^N with N >> 100 as in the DiploidModel example). And even though only a finite number of traits is alive at any time I am interested in high mutation rate regimes. There the number of different individuals alive at any time is very high.

I only came across implementations of Gillespies algorithm where there are a finite number of traits and one knows the transition rates from one trait to the other in advance (e.g. the SIR-model).

Typically I have a lot of rates to sum up to get the total rate. But unfortunately I also have to update every single rate after every single event due to the competition pressure: In the models I study individuals or die due to natural or competitive death. Every individual exceeds a small competitive pressure on every other individual. Hence if any individual exits or enters the population the death rate of every individual has to be adapted accordingly.

Luis · December 3, 2021, 10:58am

I added a SIR model to the Examples folder to compare it to the other Gillespie Algorithms around. This is the Benchmark result:

BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  25.684 μs …  2.395 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     34.669 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   38.426 μs ± 45.432 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▄▃▅▆▇▇▇█▅▂▂▂▃▁▂▂▂▁▁    ▁                                    ▂
  █████████████████████████████▇▇▇▆▅▇▇▇▇▆▆▆▆▅▅▆▅▃▃▄▅▄▄▂▄▅▄▄▄▃ █
  25.7 μs      Histogram: log(frequency) by time       102 μs <

 Memory estimate: 848 bytes, allocs estimate: 18.

Pleas try to do the analysis yourself, because I am not quite sure if I did it right.

ChrisRackauckas · December 3, 2021, 11:48am

I think this is a good clean implementation of Gillespie for small dense problems. Other implementations like the DiffEqJump ones specialize more on large problems and sparsity, while here the conditional form will only be efficient for small numbers of reactions, but would densify it.

I think one think I can see is:

https://github.com/roccminton/Gillespie_Algorithm/blob/master/MainFunctions.jl#L17

you might want to force function specialization on the higher order functions, i.e.

function run_gillespie!(time,n₀,par,execute!::F1,rates!::F2,initrates,population_history) where {F1,F2}

etc. going down. See if that helps out. But I think what will bite you in scaling is this:

https://github.com/roccminton/Gillespie_Algorithm/blob/master/Examples/SIR.jl#L22-L31

It’s not necessarily bad, but as the potential i’s get large you’re going to have a pretty nasty conditional evaluation with bad branch prediction, so it’s a choice that will accelerate small problems at the cost of large ones.

Luis · December 6, 2021, 7:39pm

Thank you for the support! I will definitely check out the DiffEqJump package!

Actually for almost all my uses the execute! function only decides between two events (resp. is), namely birth! and death!. But I belive what it slows down is the chooseevent function, when the dictionary gets big. In some applications the dictionary carries about :

https://github.com/roccminton/Gillespie_Algorithm/blob/8db5144e1cc07e3e8fc9a904ab65aabfe13300ed/MainFunctions.jl#L230-L241

But I don’t know how to accelerate this. Moreover I don’t get where all the allocations come from. Especially in the RandomMating.jl Example. There the average dictionary has 2000-3000 entries.

goerch · December 6, 2021, 7:56pm

Here is the updated JET report:

═════ 3 possible errors found ═════
┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\examples\randommating.jl:29 Main.DiploidModel.rungillespie(t, x0, model_parameter)
│┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:35 Main.DiploidModel.Gillespie.run_gillespie!(time, n₀, Main.DiploidModel.setupparameter(model_parameter, n₀), Main.DiploidModel.execute!, Main.DiploidModel.rates!, initrates, population_history)
││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:17 Main.DiploidModel.Gillespie.mainiteration!(population_history, initrates, n₀, Main.DiploidModel.Gillespie.convert(Main.DiploidModel.Gillespie.Float64, Base.getindex(time, 1)), time, par, execute!, rates!)
│││┌ @ C:\Users\Win10\.julia\packages\ProgressMeter\Vf8un\src\ProgressMeter.jl:934 meter#291 = ProgressMeter.Progress(ProgressMeter.length(iterable))
││││┌ @ C:\Users\Win10\.julia\packages\ProgressMeter\Vf8un\src\ProgressMeter.jl:94 $(QuoteNode(ProgressMeter.var"#Progress#1#2"()))(0.1, "Progress: ", :green, ProgressMeter.stderr, ProgressMeter.nothing, ProgressMeter.BarGlyphs('|', '█', _3, ' ', '|'), 0, 0, true, false, #self#, n)
│││││┌ @ C:\Users\Win10\.julia\packages\ProgressMeter\Vf8un\src\ProgressMeter.jl:94 ProgressMeter.running_ijulia_kernel()
││││││┌ @ C:\Users\Win10\.julia\packages\ProgressMeter\Vf8un\src\ProgressMeter.jl:243 Base.getproperty(ProgressMeter.Main, :IJulia)
│││││││┌ @ Base.jl:31 Base.getfield(x, f)
││││││││ variable Main.IJulia is not defined: Base.getfield(x::Module, f::Symbol)
│││││││└──────────────
│││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:62 ct = Main.DiploidModel.Gillespie.onestep!(n0, rates, ct, step, par, ex!, r!)
││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:145 Main.DiploidModel.Gillespie.iszero(i)
│││││┌ @ number.jl:42 Base.zero(x)
││││││ no matching method found for call signature (Tuple{typeof(zero), Nothing}): Base.zero(x::Nothing)
│││││└────────────────
││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\MainFunctions.jl:149 ex!(i, x_0, par)
│││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:194 Main.DiploidModel.death!(ps, par)
││││││┌ @ C:\Users\Win10\Documents\GitHub\Gillespie_Algorithm-master\DiploidModel.jl:187 Base.setindex!(ps, Main.DiploidModel.-(Base.getindex(ps, Main.DiploidModel.choosefey(ps, Base.getindex(Base.getproperty(par, :popsize), 1))), Main.DiploidModel.one(Main.DiploidModel.valtype(ps))), Main.DiploidModel.choosefey(ps, Base.getindex(Base.getproperty(par, :popsize), 1)))
│││││││┌ @ dict.jl:374 key = Base.convert(_, key0)
││││││││ no matching method found for call signature (Tuple{typeof(convert), Type{Main.DiploidModel.RandomIndividual{Int64}}, Nothing}): key = Base.convert(_::Type{Main.DiploidModel.RandomIndividual{Int64}}, key0::Nothing)
│││││││└───────────────