Gillespie simulations on GPU

[GeekyPeas]

I recently came across the package Catalyst.jl (and I am very impressed with/ excited about it!).

I need many samples (in the order of 10^5) from the Gillespie simulation of a chemical reaction network. I would like to use a GPU for this purpose. A GPU, I believe, can output in the order of 10^3 samples per batch.

So, how to do this with Catalyst.jl?

Sorry I am very new to both Julia and Catalyst.jl. In the description of Catalyst.jl it does say "High performance, GPU-parallelized, and O(1) solvers". However I couldn't find any documentation on running Catalyst.jl code on GPU.

So far this is what I am thinking:

dprob = DiscreteProblem(model,u0,tspan,p); jprob = JumpProblem(model,dprob,Direct());

are one-time operations.

It is the:

solve(jprob,SSAStepper())

that needs to be run many times on the GPU.

Using @which solve(jprob,SSAStepper()) I see this method is actually outside of the Catalyst.jl package and comes from the DiffEqBase package.

As of now, I have not figured out how code should be run on GPU in Julia. But I thought, asking may get me answers faster!

[GeekyPeas]

I found this: Parallel Ensemble Simulations from DifferentialEquations.jl.

However as I am new to Julia, can someone just tell me the steps to get Gillespie working on GPU. I would actually want to see it working before understanding why it is working.

[isaacsas]

Hi, unfortunately the GPU parallelization refers to the ODE and SDE solvers.

I'm not aware that there are any real good (exact) GPU-based Gillespie algorithms, but I haven't kept up on that literature. If you are aware of one please do let me know the reference, and perhaps we can look into adding it to JumpProcesses.jl (which would make it available to Catalyst).

[GeekyPeas]

I think this instruction should be a part of the package documentation.

[isaacsas]

We should indeed add a tutorial that shows how to use GPUs for ODE/SDEs. @ChrisRackauckas would that make sense to have in the code-optimization tutorial you were planning to work on?

[GeekyPeas]

Hi, thanks for the prompt reply! Yes a tutorial on GPUs for ODE/SDEs would be nice.

For Gillespie, I had something very obvious on mind: Essentially each trajectory or simulation of Gillespie is independent of another. So each simulation or trajectory can be simulated independently in parallel in a GPU Block.

So if one needs only a few samples of Gillespie, then using a GPU may not be very useful. However if one needed a very large number of samples, then GPU can help.

The parallelization that I am suggesting is very basic (and also perhaps very dumb/ obvious). Our group had previously implemented this on Python and we got good boosts! However I think there may be more sophisticated ways of going about this. A quick Google search yielded this: https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0046693.

[GeekyPeas]

For now, I would be happy if I can just run as many trajectories of Gillespie as my GPU would permit. Would you know how I can do that?

[isaacsas]

Which Gillespie method algorithm did you use on GPU? How big was your reaction model?

The Gillespie methods we expose are all in JumpProcesses.jl, but I doubt any of them would work on GPUs in their current implementation -- they certainly weren't written with GPU usage in mind.

[ChrisRackauckas]

The parallel ensembling tutorial for Catalyst would be kept separate from the performance optimization example because there's too much to say just on single solve performance. They are two topics with different analyses and methods, and both will be long.

As for GPUs, yes any parallel ensembling tutorial would have GPU. There's new GPU algorithms we're finalizing and we can have that happen at the same time. @utkarsh530 will be joining the Julia Lab and working on that.

For GPU Gillespie, our current representation would only be GPU compatible with pure mass action simulations. That has not been implemented yet, and other GPU Gillespie solely focus on that (like the linked paper). We could expand mass action-like handling to Hill functions as well (which I think we should do anyways to specialize the form and reduce the number of rate functions, it'll make more things faster), and then that could also be GPU compatible. But with that, I don't want to write a GPU-based Gillespie until we do that representation expansion. In theory arbitrary rate functions could be supported, but it would blow the registers if you're not careful.

[GeekyPeas]

We used partial propensity and composition rejection. We had plans to implement tau-leaping but we never got there. Our typical reaction networks had about 30 species and about 50 reactions. But we had tested our engine on synthetic examples with up to 1500 species and 4000 reactions.

[GeekyPeas]

@ChrisRackauckas I would love to contribute to this, once I get a bit more familiar with the landscape here.

For now, for my needs, pure mass action suffices. As you say, the current representation is GPU compatible with pure mass action kinetics - could you please elaborate?

[isaacsas]

We used partial propensity and composition rejection. We had plans to implement tau-leaping but we never got there. Our typical reaction networks had about 30 species and about 50 reactions. But we had tested our engine on synthetic examples with up to 1500 species and 4000 reactions.

Ah, that is great to hear. Well, as I said all our SSAs are implemented in JumpProcesses.jl, so we could work on making those work on GPUs at least in the pure mass action case. If you want to take a look at those and make PRs (or even suggestions via issues) we can see about what can be done to get some of them working on GPUs in serial.

I don't actually have real compute GPU access myself, so can't currently test / play with this.

[GeekyPeas]

@isaacsas thanks! Yes I will be looking through the Jump processes.jl code over the coming weekend.

[isaacsas]

This is more of a JumpProcesses issue so I'm going to close here, but contributions to add such functionality to JumpProcesses would be welcome!