JonnyyTorres / Galsim_JAX

This repository is intended to migrate the code from the paper "Deep Generative Models for Galaxy Image Simulations" from TensorFlow 1.x to JAX.
MIT License
1 stars 1 forks source link

Incorporation of PSF in Model #5

Open JonnyyTorres opened 1 year ago

JonnyyTorres commented 1 year ago

Hello @b-remy and @EiffL. As we discused, the last code version contains the convolution with the PSF in the loss function by using the COSMOS Dataset. Here, I'd like to present you some interesting results to know your opinion and the steps we should test. First of all, some outputs of the (Raw) Autoencoder with the SD architecture that achieved the best log-likelihood results to date:

Autoencoder_samples

The same images but convolved with the PSF:

Convolve_samples

And finally the same examples but this time taking the convolved images as the mean of a Multivariate Normal Diagonal Distribution and then sampled from this distribution

MNDist_samples

As you can appreciate the added noise in the output is quite appreciable. In general, the magnitude of the galaxies in this dataset as you know is higher, being reflected in a lower SNR compared to the HSC dataset.

This week, I have been waiting for the variance calculation for each image in the dataset. It took a considerable time, but I hope to be able to incorporate it in the loss function soon in order to observe the changes.

I remain attentive to your observations and comments :-)

JonnyyTorres commented 1 year ago

Hello @b-remy and @EiffL. Today I would like to present you some progress on the model with a better loss after trying different configurations. First the images obtained by the autoencoder:

Autoencoder_samples(1)

As you can see, I still have some diagonal stripes (with Benjamin we lowered the hypothesis that it may be due to the way the convolution was performed). Here are the images after convolution with the PSF:

Convolve_samples(1)

And finally the images after being sampled from the Diagonal Multivariate Gaussian Diagonal Distribution:

MNDist_samples(1)

As you can see, the results are starting to look a little better, although I suppose they are still susceptible to improvement, either by incorporating a total variance factor or by modifying the convolution function.

On the other hand, I share with you this gif where you can appreciate the difference between the original images and the resulting ones after going through the autoencoder :-) :

ezgif com-crop

The hyperparameters I used to obtain this result are listed below:

Key Value
act_fn "gelu"
batch_size 16
interpolation "Bicubic"
kl_reg 0.001
latent_dim 128
learning_rate 0.01
opt "adafactor"
resnet_blocks 2
steps 125,000
steps_schedule 50,000

As for the Autoencoder architecture, it is still the same as the one I usually use, that is, three symmetrical convolution layers for the Encoder and the Decoder, having two ResNet Blocks for each Layer. At the moment, the latent space dimension I am working with is 32², however, Benjamin tells me that it might be more convenient to work with a 16² dimension, test that I will perform today to see if I get some improvements in the loss values.

b-remy commented 1 year ago

I think we're good with the convolution now

difference_pred_125005_f8251ddc4554c2d1215c.png

Let's see with your best training setting @JonnyyTorres :-)