Machine Learning and Science Forum — FlowPM: Distributed TensorFlow Implementation of Cosmological N-body Solver

ML&Sci Forum

September 28, 2020
11:00am to 12:00pm
Virtual Participation

Machine Learning and Science Forum 
Date: Monday, September 28, 2020
Time: 11:00 AM - 12:00 PM Pacific Time
Location: Participate remotely using this Zoom link 

FlowPM: Distributed TensorFlow Implementation of Cosmological N-body Solver

Chirag Modi, Berkeley Center for Cosmological Physics
Abstract: Complex simulations that describe physical phenomena provide high-fidelity models in many scientific domains. However these can be poorly suited for inference and lead to challenging inverse problems. One of the most promising ways to address this and develop simulation-based inference is with differentiable simulations, and ML libraries can go a long way to this end. In this talk, I will present FlowPM, the first differentiable, distributed cosmological N-body simulation code implemented in TensorFlow. I will demonstrate how FlowPM will allow cosmologists to efficiently solve large scale inference problems in analyzing future cosmological surveys with - 1) GPU accelerations leading to more than 10x speed-up over current simulations, 2) seamless integration of FlowPM with deep learning components allowing us to push beyond traditional cosmological modeling and 3) the inbuilt differentiability of these simulations allowing us to solve inverse problems. Primary bottleneck for these N-body simulations is efficient computation of large scale forces on distributed systems, and I will show how FlowPM solves this with a novel multi-grid scheme based on multiresolution image pyramids.

The Machine Learning and Science Forum meets biweekly to discuss current applications across a wide variety of research domains in the physical sciences and beyond. Hosted by UC Berkeley Physics Professor and BIDS Faculty Affiliate Uros Seljak, these active sessions bring together domain scientists, statisticians, and computer scientists who are either developing state-of-the-art methods or are interested in applying these methods in their research. To receive email notifications about upcoming meetings, or to request more information, please contact berkeleymlforum@gmail.comAll interested members of the UC Berkeley and Berkeley Lab communities are welcome and encouraged to attend. Full details about this meeting are posted here:


Chirag Modi

Berkeley Center for Cosmological Physics