Efficient and scalable prediction of spatio-temporal stochastic gene expression in cells and tissues using graph neural networks

Abstract

The simulation of spatial stochastic models is highly computationally expensive, an issue that has severely limited our understanding of the spatial nature of gene expression. Here we devise a graph neural network based method to learn, from stochastic trajectories in a small region of space, an effective master equation for the time-dependent marginal probability distributions of mRNA and protein numbers at sub-cellular resolution for every cell in a tissue. Numerical solution of this equation leads to accurate results in a small fraction of the computation time of standard simulation methods. Moreover its predictions can be extrapolated to a spatial organisation (a cell network topology) and regions of parameter space unseen in its neural network training. The scalability and accuracy of the method suggest it is a promising approach for whole cell modelling and for detailed comparisons of stochastic models with spatial genomics data.