This repository contains code to implement Generative Adversarial Neural networks to produce images of matter distribution in the universe for different sets of cosmological parameters. Dataset consisits of N-body cosmology simulation maps obtained from PyCOLA. The pixel values represent matter density.
The aim is to build conditional GANs to produce images for different classes corresponding to different sets of cosmological parameters.
We develop a simple GAN trained on 3D images of size 64^3.
| 3D GAN: Pixel intensity | 3D GAN: Power spectrum |
|---|---|
 |
 |
Below are 2D snapshots of a set of 3D images. The input images are to the left and the GAN generated images are to the right.
| 3D GAN: Input images | 3D GAN: Generated images |
|---|---|
 |
 |
We develop a conditional GAN trained on 2D images of size
| 2D cGAN: Pixel intensity | 2D cGAN: Power spectrum |
|---|---|
 |
 |
There separate directories to run GAN and cGAN in for both 2D and 3D datasets. The Table below describes the content of various folders in this repository.
| Description | Location |
|---|---|
| 2d GAN - Main code | code/1_basic_GAN/1_main_code/ |
| 2d GAN - Analysis code | code/1_basic_GAN/2_analysis |
| 2d CGAN - code | code/3_cond_GAN/1_main_code |
| 2d CGAN - analysis | code/3_cond_GAN/2_cgan_analysis/ |
| 3d GAN | code/4_basic_3d_GAN/1_main_code/ |
| 3d CGAN - code | code/5_3d_cgan/1_main_code |
There are jupyter notebooks to build launch scripts to run the code on cori GPUs at NERSC and GPUs on SUMMIT and to perform post-run computation of metrics for different stored images. Each folder contains a jupyter notebook to quickly test the code, a folder with the full code, and a folder with analysis codes to inspect the performance of the code. Below is an example for the 2D GAN:
| Name | Description |
|---|---|
| run_scripts/launch_train_pytorch.ipynb | Notebook that launches script to run training |
| run_scripts/launch_compute_chisqr.ipynb | Notebook that launches script to run post-run metric computation |
| 1_basic_GAN/cosmogan_train.ipynb | Jupyter notebook to test GAN |
| 1_basic_GAN/1_main_code | Folder containing main training code |
| 1_basic_GAN/2_analysis/1_pytorch_analyze-results.ipynb | Notebook to analyze GAN results and view best epoch-steps |