Adds single Sersic fit parameters

[EiffL]

@mhuertascompany this PR adds parameters from single sersic fits for COSMOS galaxies. This defines the following "problem" attrs2img_cosmos_128_euclid with the following parameters:

  def hparams(self, defaults, model_hparams):
    p = defaults
    p.pixel_scale = 0.03
    p.img_len = 128
    p.example_per_shard = 1000
    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                  "attributes":  modalities.ModalityType.IDENTITY,
                  "targets": modalities.ModalityType.IDENTITY}
    p.vocab_size = {"inputs": None,
                    "attributes": None,
                    "targets": None}
    p.attributes = ['mag_auto', 'flux_radius', 'sersic_n', 'sersic_q']

(see here: https://github.com/ml4astro/galaxy2galaxy/blob/54e5c45361936c1c3e4f89f9443089442bf6ebf3/galaxy2galaxy/data_generators/cosmos.py#L266 )

Which means, the training set of the COMOS galaxies will be drawn at native 0.03 arcsec resolution on postage stamps on size 128, and the image generation will be conditioned on 'mag_auto', 'flux_radius', 'sersic_n', 'sersic_q'

To install the code from this PR:

$ pip install --user -e git+https://github.com/ml4astro/galaxy2galaxy.git@u/EiffL/additional_fit_params

I'm assuming here you already have GalSim installed, and the COSMOS dataset downloaded (see here https://github.com/GalSim-developers/GalSim/wiki/RealGalaxy%20Data)

The only thing to do now is to generate the dataset:

$ g2g-datagen --problem=attrs2img_cosmos128_euclid --data_dir=datasets/attrs2img_cosmos128_euclid

If this doesn't work, you need to add --tmp_dir=[Path to the COSMOS_25.2_training_sample ]

Now you can use the encoder/decoder that I've already trained for these images, downloadable here: https://storage.googleapis.com/g2g/modules/vae_cosmos_128/encoder.tar.gz https://storage.googleapis.com/g2g/modules/vae_cosmos_128/decoder.tar.gz

With this dataset, you can train the latent flow model this way:

$ g2g-trainer --problem=attrs2img_cosmos128_euclid --data_dir=datasets/attrs2img_cosmos128_euclid --output_dir=training/latent_maf --model=latent_maf --hparams_set=latent_flow_larger --hparams="encoder_module=vae_cosmos_128/encoder"

and then export it using g2g-exporter:

$ g2g-exporter --problem=attrs2img_cosmos128_euclid --data_dir=datasets/attrs2img_cosmos128_euclid --output_dir=training/latent_maf --model=latent_maf --hparams_set=latent_flow_larger --hparams="encoder_module=vae_cosmos_128/encoder" --export_dir=modules/latent_maf

This exports the latent space model as a TF Hub module. A final step to combine the latent space model and the decoder into a single generator model:

$ python galaxy2galaxy/bin/concatenate_models.py --decoder_module=vae_cosmos_128/decoder --flow_module=modules/latent_maf --export_dir=modules/flow_vae_cosmos_128

And you are done, after that you can use GalSim Hub to generate the images, see this presentation from January: https://slides.com/eiffl/galaxy_morphology#/3

import galsim
import galsim_hub

STAMP_SIZE = 128 # Desired postage stamp size, can be different from training size
PIXEL_SCALE = 0.03 # Desired pixel scale, can be different from training pixel scale

# Loads COSMOS sample
cosmos_cat = galsim.COSMOSCatalog()

# Loads model for flow-vae
model = galsim_hub.GenerativeGalaxyModel('modules/flow_vae_cosmos_128/generator')

# Noise model
cosmos_noise = galsim.getCOSMOSNoise()

figure(figsize=(17.5,6))
for i in range(9):
    im_real = galsim.ImageF(STAMP_SIZE, STAMP_SIZE, scale=PIXEL_SCALE)
    im_sim = galsim.ImageF(STAMP_SIZE, STAMP_SIZE, scale=PIXEL_SCALE)
    im_sim_n = galsim.ImageF(STAMP_SIZE, STAMP_SIZE, scale=PIXEL_SCALE)

    # Drawing galaxy from cosmos
    gal = cosmos_cat.makeGalaxy(i, gal_type='real')
    psf = gal.original_psf
    real = galsim.Convolve(psf, gal)
    real.drawImage(im_real, method='no_pixel');

    # Drawing galaxy from generative model
    params = cosmos_cat.param_cat[cosmos_cat.orig_index][['mag_auto', 'flux_radius', 'sersic_n', 'sersic_q']][i:i+1]
    sim_gal = model.sample(params)
    sim = galsim.Convolve(psf, sim_gal)
    sim.drawImage(im_sim, method='no_pixel');
    sim.drawImage(im_sim_n, method='no_pixel');
    im_sim_n.addNoise(cosmos_noise)

    subplot(3,9,i+1)
    imshow(arcsinh(50*im_real.array), cmap='gray')
    axis('off')

    subplot(3,9,i+1+9)
    imshow(arcsinh(50*im_sim_n.array), cmap='gray')   
    axis('off')

    subplot(3,9,i+1+18)
    imshow(arcsinh(50*im_sim.array), cmap='gray')   
    axis('off')

plt.subplots_adjust(wspace=0, hspace=0)

[EiffL]

@nesar @cguilloteau Until I get around to writing an actual documentation, this thread can also be of interest to you, it shows how to train a latent space conditional MAF

[EiffL]

Ok, I'm gonna go ahead and merge this as it's already used by @Hbretonniere