FlowPM

Particle Mesh Simulation in TensorFlow, based on fastpm-python simulations

Try me out:

To install:

$ pip install flowpm

For a minimal working example of FlowPM, see this notebook. The steps are as follows:

import tensorflow as tf
import numpy as np
import flowpm

cosmo = flowpm.cosmology.Planck15()
stages = np.linspace(0.1, 1.0, 10, endpoint=True)

initial_conditions = flowpm.linear_field(32,          # size of the cube
                                         100,         # Physical size of the cube
                                         ipklin,      # Initial power spectrum
                                         batch_size=16)

# Sample particles
state = flowpm.lpt_init(cosmo, initial_conditions, a0=0.1)   

# Evolve particles down to z=0
final_state = flowpm.nbody(cosmo, state, stages, 32)         

# Retrieve final density field
final_field = flowpm.cic_paint(tf.zeros_like(initial_conditions), final_state[0])

Mesh TensorFlow implementation

FlowPM provides a Mesh TensorFlow implementation of FastPM, for running distributed simulations across very large supercomputers.

Instructions for GPU clusters

We rely on a customized Mesh TensorFlow backend based on Horovod to distribute computations on GPU clusters through the high performance NCCL library.

To install the necessary dependencies, you first need to be in an environment providing:

• TensorFlow 2.1 or above
• NCCL 2.8 or above You can then install Horovod and Mesh TensorFlow with:

$ pip install git+https://github.com/horovod/horovod.git
$ pip install git+https://github.com/DifferentiableUniverseInitiative/mesh@hvd_max

TPU setup

To run FlowPM on Google TPUs here is the procedure

• Step 1: Setting up a cloud TPU in the desired zone, do from the GCP console:

$ gcloud config set compute/region europe-west4
$ gcloud config set compute/zone europe-west4-a
$ ctpu up --name=flowpm --tpu-size=v3-32

• Step 2: Installing dependencies and FlowPM:

$ git clone https://github.com/DifferentiableUniverseInitiative/flowpm.git
$ cd flowpm
$ git checkout mesh
$ pip3 install --user mesh-tensorflow
$ pip3 install --user -e .

It's so easy, it's almost criminal.

Notes on using and profiling for TPUs

There a few things to keep in mind when using TPUs, in particular, the section on Excessive tensor padding from this document: https://cloud.google.com/tpu/docs/troubleshooting

See the README in the script folder for more info on how to profile