Kavli_UCSC_2019

The code is used for the project of Kavli summer program in Astrophysics in 2019.

What do we want to do?

In this project, we are exploring the usage of VQ-VAE (1) (2) to two ideas:

• Emulating galaxy images (without noise/psf) to adapt to other surveys using generative models (e.g. PixelCNN)
• Unsupervised machine learning (clustering) to explore the morphology of galaxies (has been moved to another repo)

We modified the VQ-VAE example (Source code of VQ-VAE) from Sonnet library.

Preliminary results

VQ-VAE code on CANDELS data

$python image_to_dict.py

This code is used to generate a dictionary to store images.

$python vqvae.py

I run 100000 epochs on one GPU which takes ~45 mins (~5400 images in training).
The data is from CANDELS (GOODS, h band).

(Change line 155: Train=False, to use pretrained VQ-VAE model.)

Add noise/PSF layers

$python image_to_dict_noise+psf.py

This code is used to generate a dictionary to store images when you have PSF file. The noise map is generated by Gaussian noise using the mean and variance calculated at 4 corners in each image.

$python vqvae_noise+psf.py

I added noise and PSF layers in the decoder and run 100000 epochs for this example on one GPU which takes ~ 4hrs (~5400 images in training). The data is from CANDELS (GOODS, h band), and the PSF image is from 3DHST (GOODS-S WFC3 PSFs, F160W).

Connect with Gated PixelCNN

In this section, I trained Gated PixelCNN on vector quantized feature map. After implementing VQVAE, we will have two outputs from VQVAE which are the inputs for PixelCNN:

• codebook (e): shape=(K,D), where K is the number of code and D is the dimension of the code.
• encodings indice: shape=( , latent_size, latent_size, dtype=int32), where 'latent_size' means the size of latent map.

I followed the code from hiwonjoon which tried to reproduce the results of the paper - VQ-VAE (2). They applied the code of Gated PixelCNN on the vector quantized feature map extracted from VQ-VAE.

I trained PixelCNN (unconditional) 100000 epoches on both CANDELS (84x84) (~7000 images) and sample the prior to generate new feature map. I then use a pre-trained decoder to reconstruct new images from generated feature maps.

• CANDELS result (deconvolution)

To-do list

• Adapt the VQ-VAE code to astronomical images (e.g. CANDELS)
• Add noise/PSF layers before output and retrieve the reconstructed images from the layer before noise/PSF layers
• Connect with PixelCNN to generate random galaxy images
• SDSS data
• Conditional to what labels? --> we condition on T-Type at the moment.
• t-SNE to explore the meaning of quantized vector and compare with galaxy properties
• Examination