Simulating the Universe in TensorFlow

Guest post by , from Berkeley Center for Cosmological Physics, CosmoStat Laboratory, and the Lawrence Berkeley National Laboratory¹ of the large scale structure of the Universe are fundamental tools used by cosmologists to make sense of the vast amount of data collected by cosmological surveys. These simulations are typically extremely computationally expensive and are usually run offline on massive supercomputers. However, what if we could make these simulations extremely fast and integrate them with machine learning components in a single unified framework? This is what a new N-body cosmological simulation code, , a pure TensorFlow implementation of cosmological N-body simulations. We provide a (BCCP, UC Berkeley), we are interested in going backwards in time and reconstructing the with comes in.

Mesh TensorFlow framework allows us to easily describe our simulation in terms of distributed tensors, keeping track behind the scenes of distributed gradients and memory communications between devices. By writing our In addition to enabling large simulations, a model parallelism framework also allows us to speed up intermediate size simulations by splitting the computation across multiple processors. This is demonstrated in the following Figure 4 where we show that on average, FlowPM simulations are 40x faster than the current differentiable python simulations, FastPM.

Figure 4 : We compare the time scaling with number of processors for 1 step in 2563 grid PM simulation in FastPM (CPU based python code run on Cori Haswell cores) & FlowPM (GPU based Mesh TensorFlow code run on Cori GPUs)

Outlook

Numerical simulations of our Universe have formed the backbone of the large scale structure cosmology for more than three decades. With FlowPM, we are taking the first steps to integrate these simulations with deep learning components in a single unified framework while maintaining the exact physical understanding of the underlying phenomenon. In cosmology, this combination has opened doors to developing novel analytic tools as well as push modeling into regimes that was hitherto intractable. These are areas of active research, made increasingly urgent with the next generation of cosmological surveys, that will observe tens of millions of objects in the Universe coming online at the turn of the decade. This confluence of physical modeling and machine learning has largely been made possible due to the model parallelism framework of Mesh TensorFlow, and we hope that the component analytic and computing tools developed with FlowPM will also benefit large scale scientific applications in other disciplines beyond cosmology.

We would like to earnestly acknowledge the support of our colleagues at NERSC - Wahid Bhimji, Steve Farrell, Peter Harrington, Prabhat and at Google - Niki Parmar, Thiru Palanisamy, Noam Shazeer, Youlong Cheng, Zak Stone as well as others who have pointed us to relevant resources, actively discussed ways to optimize and improve these simulations and provided useful feedback.

References:

1. FastPM (underlying PM scheme for FlowPM)-
   
2. Dai et al.:
   
3. Mesh Tensorflow :
   
4. Parallel FlowPM code with MeshTF: https://github.com/modichirag/flowpm/tree/mesh