Skip to main content

Particle Mesh Simulation in TensorFlow

Project description

flowpm Build Status PyPI version Open In Colab arXiv:2010.11847 youtube PEP8

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

Try me out: Open In Colab

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

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(initial_conditions, a0=0.1)   

# Evolve particles down to z=0
final_state = flowpm.nbody(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.

Here are the instructions for installing and running on Cori-GPU. More info about this machine here: https://docs-dev.nersc.gov/cgpu/

  1. Login to a cori-gpu node to prepare the environment:
$ module add esslurm
$ salloc -C gpu -N 1 -t 30 -c 10 --gres=gpu:1 -A m1759
  1. First install dependencies
$ module purge && module load gcc/7.3.0 python3 cuda/10.1.243
$ pip install --user tensorflow==2.1
$ pip install --user mesh-tensorflow

NOTE: we are installing our own tensorflow 2.1 version until a module is available at NERSC

  1. Install the Mesh TensorFlow branch of FlowPM
$ git clone https://github.com/modichirag/flowpm.git
$ cd flowpm
$ git checkout mesh
$ pip install --user -e .
  1. To run the demo comparing the distributed computation to single GPU:
$ cd examples
$ sbatch lpt_job.sh

This will generate a plot comparison.png showing from a set of initial conditions, the result of a single LPT step on single GPU TensorFlow vs Mesh TensorFlow.

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/modichirag/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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for flowpm, version 0.1.1
Filename, size File type Python version Upload date Hashes
Filename, size flowpm-0.1.1.tar.gz (1.6 MB) File type Source Python version None Upload date Hashes View

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page