Use galsim as realistic source galaxy generator

.gitignore

@@ -31,4 +31,8 @@ docs/_build/*
 # Temp
 *Untitled.ipynb
 *~
-*#
+*#
+
+# macos cache files
+.DS_Store
+**/.DS_Store

baobab/bnn_priors/__init__.py

@@ -3,4 +3,5 @@
 from .cov_bnn_prior import CovBNNPrior
 from .empirical_bnn_prior import EmpiricalBNNPrior
 from .base_cosmo_bnn_prior import BaseCosmoBNNPrior
-from .diagonal_cosmo_bnn_prior import DiagonalCosmoBNNPrior
+from .diagonal_cosmo_bnn_prior import DiagonalCosmoBNNPrior
+from .real_source_prior import RealSourcePrior

baobab/bnn_priors/base_bnn_prior.py

@@ -5,8 +5,9 @@ class BaseBNNPrior(ABC):
     """Abstract base class equipped with PDF evaluation and sampling utility functions for various lens/source macromodels
 
     """
-    def __init__(self, bnn_omega, components):
+    def __init__(self, bnn_omega, components, external=None):
         self.components = components
+        self.external = external
         for comp in bnn_omega:
             setattr(self, comp, bnn_omega[comp])
         self._set_required_parameters()
@@ -60,7 +61,8 @@ def _set_required_parameters(self):
                           SHEAR_GAMMA_PSI=['gamma_ext', 'psi_ext'],
                           SERSIC_ELLIPSE=['magnitude', 'center_x', 'center_y', 'n_sersic', 'R_sersic', 'e1', 'e2'],
                           LENSED_POSITION=['magnitude'],
-                          SOURCE_POSITION=['ra_source', 'dec_source', 'magnitude'],)
+                          SOURCE_POSITION=['ra_source', 'dec_source', 'magnitude'],
+                          GALSIM=['catalog_index', 'galsim_scale', 'galsim_angle', 'galsim_center_x', 'galsim_center_y'])  # TODO : add magnitude
         setattr(self, 'params', params)
 
     def _raise_config_error(self, missing_key, parent_config_key, bnn_prior_class):
@@ -92,4 +94,4 @@ def sample(self):
         Overridden by subclasses.
 
         """
-        return NotImplemented
+        return NotImplemented

baobab/bnn_priors/real_source_prior.py

@@ -0,0 +1,151 @@
+from addict import Dict
+import lenstronomy.Util.param_util as param_util
+from lenstronomy.LensModel.lens_model import LensModel
+from lenstronomy.LensModel.lens_model_extensions import LensModelExtensions
+from .base_bnn_prior import BaseBNNPrior
+
+class RealSourcePrior(BaseBNNPrior):
+    """BNN prior with independent parameters
+
+    Note
+    ----
+    first test for pixelated realistic source training set
+
+    """
+    def __init__(self, bnn_omega, components, external):
+        """
+        Note
+        ----
+        The dictionary attributes are copies of the config corresponding to each component.
+        The number of attributes depends on the number of components.
+
+        Attributes
+        ----------
+        components : list
+            list of components, e.g. `lens_mass`
+        lens_mass : dict
+            profile type and parameters of the lens mass
+        src_light : dict
+            profile type and parameters of the source light
+
+        """
+        BaseBNNPrior.__init__(self, bnn_omega, components, external=external)
+        self.params_to_exclude = []
+        self.set_params_list(self.params_to_exclude)
+        self.set_comps_qphi_to_e1e2()
+        if 'src_light' in self.components and bnn_omega.src_light.profile in ['GALSIM']:  # can add other here
+            self._src_light_is_pixel = True
+        else:
+            self._src_light_is_pixel = False
+        lens_model_list = [bnn_omega.lens_mass.profile]
+        self._lensModel = LensModel(lens_model_list)
+        self._lensModelExt = LensModelExtensions(self._lensModel)
+
+    def sample(self):
+        """Gets kwargs of sampled parameters to be passed to lenstronomy
+
+        Returns
+        -------
+        dict
+            dictionary of config-specified components (e.g. lens mass), itself
+            a dictionary of sampled parameters corresponding to the config-specified
+            profile of that component
+
+        """
+        # Initialize nested dictionary of kwargs
+        kwargs = Dict()
+
+        # Realize samples
+        for comp, param_name in self.params_to_realize:
+            hyperparams = getattr(self, comp)[param_name].copy()
+            param_value = self.sample_param(hyperparams)
+            kwargs[comp][param_name] = param_value
+
+        # We want the source to be close to the caustics
+        margin = 1  # units of pixel size
+        kwargs = self._ensure_center_in_caustics('src_light', kwargs, margin=margin, verbose=False)
+
+        # Ext shear is defined wrt the lens center
+        kwargs['external_shear']['ra_0'] = kwargs['lens_mass']['center_x']
+        kwargs['external_shear']['dec_0'] = kwargs['lens_mass']['center_y']
+
+        # Source pos is defined wrt the lens pos
+        if 'galsim_center_x' in kwargs['src_light']:
+            kwargs['src_light']['galsim_center_x'] += kwargs['lens_mass']['center_x']
+            kwargs['src_light']['galsim_center_y'] += kwargs['lens_mass']['center_y']
+        else:
+            kwargs['src_light']['center_x'] += kwargs['lens_mass']['center_x']
+            kwargs['src_light']['center_y'] += kwargs['lens_mass']['center_y']
+
+        if 'lens_light' in self.components:
+            # Lens light shares center with lens mass
+            kwargs['lens_light']['center_x'] = kwargs['lens_mass']['center_x']
+            kwargs['lens_light']['center_y'] = kwargs['lens_mass']['center_y']
+
+        # In case of pixelated light such as galsim galaxies, translate to 'interpol' profile of lenstronomy
+        self.original_sample = kwargs.copy()  # save for public access
+        if self._src_light_is_pixel:
+            kwargs_interpol = self._kwargs_pixel2interpol(kwargs['src_light'])
+            kwargs['src_light'] = kwargs_interpol
+
+        # Convert any q, phi into e1, e2 as required by lenstronomy
+        for comp in self.comps_qphi_to_e1e2: # e.g. 'lens_mass'
+            q = kwargs[comp].pop('q')
+            phi = kwargs[comp].pop('phi')
+            e1, e2 = param_util.phi_q2_ellipticity(phi, q)
+            kwargs[comp]['e1'] = e1
+            kwargs[comp]['e2'] = e2
+
+        return kwargs
+
+    def setup_pixel_profiles(self, image, instruments, numerics):
+        self._pixel_image_size  = image.num_pix
+        self._pixel_pixel_scale = instruments.pixel_scale
+        # we want to take into account the supersampling factor, ensuring the best resolution on the pixel grid
+        self._pixel_supersampling_factor = numerics.get('supersampling_factor', 1)
+
+    def _kwargs_pixel2interpol(self, kwargs_pixel):
+        """Converts sampled parameters destined to pixelated realistic light profiles
+        into parameters supported by the interpolation scheme of lenstronomy ('INTERPOL' profile)
+        """
+        if not hasattr(self, '_pixel_image_size'):
+            raise ValueError("Image size, pixel scale, etc. not provided, use self._pixel_image_size() to do so")
+        kwargs_setup = self.external['src_light']
+        if self.src_light.profile == 'GALSIM':
+            from baobab.sim_utils import galsim_utils
+            kwargs_interpol = galsim_utils.kwargs_galsim2interpol(self._pixel_image_size, self._pixel_pixel_scale, 
+                                                                  self._pixel_supersampling_factor, kwargs_setup, kwargs_pixel)
+        else:
+            raise NotImplementedError("Pixelated light profiles other than 'GALSIM' not implemented")
+        return kwargs_interpol
+
+    def _ensure_center_in_caustics(self, comp, kwargs, margin=0, verbose=False):
+        """
+        Re-realize (galsim_center_x, galsim_center_y) of source light if it is too far from the caustics.
+        For now the simple criterion is the rectangle defined by min/max coordinates of caustics.
+        """
+        if comp != 'src_light':
+            return kwargs
+        hyperparams_x = getattr(self, comp)['galsim_center_x'].copy()
+        hyperparams_y = getattr(self, comp)['galsim_center_y'].copy()
+        current_value_x = kwargs[comp]['galsim_center_x']
+        current_value_y = kwargs[comp]['galsim_center_y']
+
+        # compute caustics by ray-shooting critical lines
+        kwargs_lens = [kwargs['lens_mass']]
+        x_crit_list, y_crit_list = self._lensModelExt.critical_curve_tiling(kwargs_lens, compute_window=5, 
+                                                                            start_scale=0.5, max_order=10)
+        x_caustic_list, y_caustic_list = self._lensModel.ray_shooting(x_crit_list, y_crit_list, kwargs_lens) 
+
+        # sample again if centroids is not inside the rectangle defined by caustics
+        param_value_x, param_value_y = current_value_x, current_value_y
+        while not (min(x_caustic_list)-margin <= param_value_x and param_value_x <= max(x_caustic_list)+margin and
+                   min(y_caustic_list)-margin <= param_value_y and param_value_y <= max(y_caustic_list)+margin):
+            param_value_x = self.sample_param(hyperparams_x)
+            param_value_y = self.sample_param(hyperparams_y)
+            kwargs[comp]['galsim_center_x'] = param_value_x
+            kwargs[comp]['galsim_center_y'] = param_value_y
+            if verbose:
+                print("Center offset ({}, {}) replaced by ({}, {})".format(current_value_x, current_value_y, 
+                                                                           param_value_x, param_value_y))
+        return kwargs

baobab/configs/__init__.py

@@ -5,4 +5,5 @@
 from . import tdlmc_diagonal_cosmo_config
 from . import gamma_diagonal_config
 from . import gamma_cov_config
-from . import gamma_empirical_config
+from . import gamma_empirical_config
+from . import real_source_config

baobab/configs/real_source_config.py

@@ -0,0 +1,162 @@
+import os
+import numpy as np
+from addict import Dict
+
+base_path = '/Users/aymericg/Documents/EPFL/PhD_LASTRO/Code/Forked/baobab_forked/demo_test'
+
+cfg = Dict()
+cfg.name = 'real_source'
+cfg.bnn_prior_class = 'RealSourcePrior'
+cfg.out_dir = os.path.join(base_path, 'galsim_test')
+cfg.seed = 213 # random seed
+cfg.n_data = 50 # number of images to generate
+cfg.train_vs_val = 'train'
+cfg.components = ['lens_mass', 'external_shear', 'src_light']  #, 'lens_light']
+cfg.checkpoint_interval = 1
+
+cfg.selection = dict(
+                 magnification=dict(
+                                    min=2.0
+                                    ),
+                 initial=[] #["lambda x: x['lens_mass']['theta_E'] > 0.5",]
+                 )
+
+cfg.instrument = dict(
+              pixel_scale=0.08, # scale (in arcseonds) of pixels
+              ccd_gain=4.5, # electrons/ADU (analog-to-digital unit). A gain of 8 means that the camera digitizes the CCD signal so that each ADU corresponds to 8 photoelectrons.
+              )
+
+cfg.bandpass = dict(
+                magnitude_zero_point=25.9463, # (effectively, the throuput) magnitude in which 1 count per second per arcsecond square is registered (in ADUs)
+                )
+
+cfg.observation = dict(
+                  exposure_time=100.0, # exposure time per image (in seconds)
+                  )
+
+cfg.psf = dict(
+           type='PIXEL', # string, type of PSF ('GAUSSIAN' and 'PIXEL' supported)
+           kernel_size=91, # dimension of provided PSF kernel, only valid when profile='PIXEL'
+           which_psf_maps=101, # None if rotate among all available PSF maps, else seed number of the map to generate all images with that map
+           )
+
+cfg.numerics = dict(
+                supersampling_factor=3  # prevent 5th image at the center due to numerical issues with cuspy mass profiles
+                )
+
+cfg.image = dict(
+             num_pix=99, # cutout pixel size
+             inverse=False, # if True, coord sys is ra to the left, if False, to the right 
+             )
+
+# put here some general (fixed) settings that are passed to the mass/light profiles (when supported)
+cfg.external = dict(
+                    src_light = dict(
+                            # the following parameters are meant to be used with 'GALSIM' profile
+                            galaxy_type='real',
+                            psf_type='gaussian',
+                            psf_gaussian_fwhm=0.12,
+                            # psf_pixel_size=0.074,  # only with psf_type='real'
+                            # psf_size=49,  # only with psf_type='real'
+                            no_convolution=False,
+                            catalog_dir='/Users/aymericg/Documents/EPFL/PhD_LASTRO/Code/divers/GalSim-releases-2.2/examples/data',
+                            catalog_name='real_galaxy_catalog_23.5_example.fits',
+                            draw_image_method='auto',
+                        )
+                    )
+
+cfg.bnn_omega = dict(
+                 lens_mass = dict(
+                                 profile='SPEMD',
+                                 center_x = dict(
+                                          dist='normal',
+                                          mu=0.0,
+                                          sigma=1.e-6,
+                                          ),
+                                 center_y = dict(
+                                          dist='normal',
+                                          mu=0.0,
+                                          sigma=1.e-6,
+                                          ),
+                                 gamma = dict(
+                                              dist='normal',
+                                              mu=2,
+                                              sigma=0.05,
+                                              lower=1.6,
+                                              upper=2.4
+                                              ),
+                                 theta_E = dict(
+                                                dist='normal',
+                                                mu=1.2,
+                                                sigma=0.05,
+                                                lower=0.8,
+                                                upper=1.6
+                                                ),
+                                 e1 = dict(
+                                          dist='beta',
+                                          a=4.0,
+                                          b=4.0,
+                                          lower=-0.5,
+                                          upper=0.5),
+                                e2 = dict(
+                                          dist='beta',
+                                          a=4.0,
+                                          b=4.0,
+                                          lower=-0.5,
+                                          upper=0.5),
+                                 ),
+
+                 external_shear = dict(
+                                       profile='SHEAR_GAMMA_PSI',
+                                       gamma_ext = dict(
+                                                         dist='normal',
+                                                         mu=0, # See overleaf doc
+                                                         sigma=0.05,
+                                                         ),
+                                       psi_ext = dict(
+                                                     dist='normal',
+                                                     mu=0,
+                                                     sigma=0.05,
+                                                     lower=-np.pi,
+                                                     upper=np.pi,
+                                                     )
+                                       ),
+
+                 src_light = dict(
+                                profile='GALSIM',
+
+                                catalog_index = dict(
+                                             dist='uniform',
+                                             lower=0,
+                                             upper=95,
+                                             ),
+                                magnitude = dict(
+                                             dist='normal',
+                                             mu=20.407,
+                                             sigma=1,
+                                             ),
+                                galsim_scale = dict(
+                                            dist='normal',
+                                            mu=0.8,
+                                            sigma=0.05,
+                                            lower=0.6,
+                                            upper=1,
+                                            ),
+                                galsim_angle = dict(
+                                             dist='uniform',
+                                             lower=-np.pi,
+                                             upper=np.pi,
+                                             ),
+                                galsim_center_x = dict(
+                                                dist='uniform',
+                                                lower=-2,
+                                                upper=2,
+                                                ),
+                                galsim_center_y = dict(
+                                                dist='uniform',
+                                                lower=-2,
+                                                upper=2,
+                                                ),
+                                ),
+
+                 )

baobab/distributions/distributions.py

@@ -32,7 +32,7 @@ def sample_uniform(lower, upper):
     sample = lower + (upper - lower)*u
     return sample
 
-def eval_uniform_pdf(eval_at, lower, upper):
+def eval_uniform_pdf(eval_at, lower, upper, force_integer=False):
     """Evaluate the uniform PDF
 
     See `sample_uniform` for parameter definitions.

baobab/generate.py

@@ -171,4 +171,4 @@ def main():
     pbar.close()
 
 if __name__ == '__main__':
-    main()
+    main()