Esempio n. 1
0
def _arc_model(numpix, kwargs_band, kwargs_model, kwargs_lens, kwargs_source_mag=None, kwargs_lens_light_mag=None,
               kwargs_ps_mag=None, kwargs_numerics={}):
    """
    routine to simulate a lensing arc, wrapper around lenstronomy SimulationAPI module

    :param numpix: number of pixels per axis
    :param kwargs_band: keyword arguments specifying the observation to be simulated according to lenstronomy.SimulationAPI
    :param kwargs_model: keyword arguments of model configurations. All possibilities available at lenstronom.Util.class_creator
    :param kwargs_lens: list of lens model keyword arguments
    :param kwargs_source_mag: list of extended source model keyword arguments
    :param kwargs_lens_light_mag: list of lens light keyword arguments
    :param kwargs_ps_mag: list of point source keyword arguments
    :param kwargs_numerics: keyword arguments describing the numerical setting of lenstronomy as outlined in lenstronomy.ImSim.Numerics
    :return: 2d numpy array
    """

    simAPI = SimAPI(numpix=numpix, kwargs_single_band=kwargs_band, kwargs_model=kwargs_model,
                    kwargs_numerics=kwargs_numerics)

    imSim = simAPI.image_model_class
    kwargs_lens_light, kwargs_source, kwargs_ps = simAPI.magnitude2amplitude(kwargs_source_mag=kwargs_source_mag,
                                                                             kwargs_lens_light_mag=kwargs_lens_light_mag,
                                                                             kwargs_ps_mag=kwargs_ps_mag)
    image = imSim.image(kwargs_lens, kwargs_source=kwargs_source, kwargs_lens_light=kwargs_lens_light,
                        kwargs_ps=kwargs_ps)
    return image
Esempio n. 2
0
class TestModelAPI(object):

    def setup(self):

        numpix = 10
        instrument_name = 'LSST'
        observation_name = 'LSST_g_band'
        kwargs_single_band = constructor.observation_constructor(instrument_name=instrument_name,
                                                          observation_name=observation_name)
        kwargs_single_band['data_count_unit'] = 'e-'
        kwargs_model = {'lens_model_list': ['SIS'], 'z_lens': None, 'z_source': None, 'lens_redshift_list': None,
                        'source_light_model_list': ['GAUSSIAN'],
                        'lens_light_model_list': ['SERSIC'], 'point_source_model_list':['UNLENSED'],
                        'source_redshift_list': None}
        kwargs_numerics = {'supersampling_factor': 2}

        self.api = SimAPI(numpix, kwargs_single_band, kwargs_model)

    def test_image_model_class(self):
        model = self.api.image_model_class()
        assert model.LensModel.lens_model_list[0] == 'SIS'

    def test_magnitude2amplitude(self):
        kwargs_lens_light_mag = [{'magnitude': 28, 'R_sersic': 1., 'n_sersic': 2, 'center_x': 0, 'center_y': 0}]
        kwargs_source_mag = [{'magnitude': 30, 'sigma': 0.3, 'center_x': 0, 'center_y': 0}]
        kwargs_ps_mag = [{'magnitude': [30], 'ra_image': [0], 'dec_image': [0]}]
        kwargs_lens_light, kwargs_source, kwargs_ps = self.api.magnitude2amplitude(kwargs_lens_light_mag, kwargs_source_mag,
                                                                        kwargs_ps_mag)

        npt.assert_almost_equal(kwargs_source[0]['amp'], 1, decimal=5)
        npt.assert_almost_equal(kwargs_ps[0]['point_amp'][0], 1, decimal=5)
        npt.assert_almost_equal(kwargs_lens_light[0]['amp'], 0.38680586575451237, decimal=5)
Esempio n. 3
0
def sim_image(numpix,
              kwargs_band,
              kwargs_model,
              kwargs_params,
              kwargs_numerics={},
              with_noise=True):
    """
    simulates an image based on chosen model and data settings, effectively makes use of lenstronomy.SimulationAPI

    :param numpix: number of pixels per axis
    :param kwargs_band: keyword arguments specifying the observation to be simulated according to lenstronomy.SimulationAPI
    :param kwargs_model: keyword arguments of model configurations. All possibilities available at lenstronom.Util.class_creator
    :param kwargs_params: keyword arguments of the different model components. Supports 'kwargs_lens', 'kwargs_source_mag',
    'kwargs_lens_light_mag', 'kwargs_ps_mag'
    :param kwargs_numerics: keyword arguments describing the numerical setting of lenstronomy as outlined in lenstronomy.ImSim.Numerics
    :return: 2d numpy array
    """

    sim = SimAPI(numpix=numpix,
                 kwargs_single_band=kwargs_band,
                 kwargs_model=kwargs_model,
                 kwargs_numerics=kwargs_numerics)
    kwargs_lens_light_mag = kwargs_params.get('kwargs_lens_light_mag', None)
    kwargs_source_mag = kwargs_params.get('kwargs_source_mag', None)
    kwargs_ps_mag = kwargs_params.get('kwargs_ps_mag', None)
    kwargs_lens_light, kwargs_source, kwargs_ps = sim.magnitude2amplitude(
        kwargs_lens_light_mag, kwargs_source_mag, kwargs_ps_mag)
    imSim = sim.image_model_class
    image = imSim.image(kwargs_params['kwargs_lens'], kwargs_source,
                        kwargs_lens_light, kwargs_ps)
    if with_noise is True:
        image += sim.noise_for_model(model=image)
    return image
Esempio n. 4
0
    def __init__(self,
                 source_x,
                 source_y,
                 variability_func,
                 numpix,
                 kwargs_single_band,
                 kwargs_model,
                 kwargs_numerics,
                 kwargs_lens,
                 kwargs_source_mag=None,
                 kwargs_lens_light_mag=None,
                 kwargs_ps_mag=None):
        """

        :param source_x: RA of source position
        :param source_y: DEC of source position
        :param variability_func: function that returns a brightness (in magnitude) as a function of time t
        :param numpix: number of pixels per axis
        :param kwargs_single_band:
        :param kwargs_model:
        :param kwargs_numerics:
        :param kwargs_lens:
        :param kwargs_source_mag:
        :param kwargs_lens_light_mag:
        :param kwargs_ps_mag:
        """

        # create background SimAPI class instance
        sim_api_bkg = SimAPI(numpix, kwargs_single_band, kwargs_model)
        image_model_bkg = sim_api_bkg.image_model_class(kwargs_numerics)
        kwargs_lens_light, kwargs_source, kwargs_ps = sim_api_bkg.magnitude2amplitude(
            kwargs_lens_light_mag, kwargs_source_mag, kwargs_ps_mag)
        self._image_bkg = image_model_bkg.image(kwargs_lens, kwargs_source,
                                                kwargs_lens_light, kwargs_ps)
        # compute image positions of point source
        x_center, y_center = sim_api_bkg.data_class.center
        search_window = np.max(sim_api_bkg.data_class.width)
        lensModel = image_model_bkg.LensModel
        solver = LensEquationSolver(lensModel=lensModel)
        image_x, image_y = solver.image_position_from_source(
            source_x,
            source_y,
            kwargs_lens,
            min_distance=0.1,
            search_window=search_window,
            precision_limit=10**(-10),
            num_iter_max=100,
            arrival_time_sort=True,
            x_center=x_center,
            y_center=y_center)
        mag = lensModel.magnification(image_x, image_y, kwargs_lens)
        dt_days = lensModel.arrival_time(image_x, image_y, kwargs_lens)
        dt_days -= np.min(
            dt_days
        )  # shift the arrival times such that the first image arrives at t=0 and the other
        #  times at t>=0
        # add image plane source model
        kwargs_model_ps = {'point_source_model_list': ['LENSED_POSITION']}
        self.sim_api_ps = SimAPI(numpix, kwargs_single_band, kwargs_model_ps)
        self._image_model_ps = self.sim_api_ps.image_model_class(
            kwargs_numerics)
        self._kwargs_lens = kwargs_lens
        self._dt_days = dt_days
        self._mag = mag
        self._image_x, self._image_y = image_x, image_y
        self._variability_func = variability_func
def make_lens_sys(lensmodel, src_camera_kwargs, source_info, kwargs_band_src,
                  lens_info):
    #Model
    kwargs_model_postit = {
        'lens_model_list': ['SIE'],
        'source_light_model_list': ['INTERPOL']
    }
    kwargs_lens = lensmodel['kwargs_lens_list']  # SIE model

    #data
    numpix = len(source_info['image']) * source_info['HR_factor']
    kwargs_source_mag = [{
        'magnitude':
        source_info['magnitude'],
        'image':
        source_info['image'],
        'scale':
        source_info['deltapix'] / source_info['HR_factor'],
        'phi_G':
        0,
        'center_x':
        lensmodel['source_shift'][0],
        'center_y':
        lensmodel['source_shift'][1]
    }]  #phi_G is to rotate, centers are to shift in arcsecs
    sim = SimAPI(numpix=numpix,
                 kwargs_single_band=kwargs_band_src,
                 kwargs_model=kwargs_model_postit)
    kwargs_numerics = {'supersampling_factor': source_info['HR_factor']}
    imSim = sim.image_model_class(kwargs_numerics)
    _, kwargs_source, _ = sim.magnitude2amplitude(
        kwargs_source_mag=kwargs_source_mag)

    #simulation
    image_HD = imSim.image(kwargs_lens=kwargs_lens,
                           kwargs_source=kwargs_source)

    mag_LS = -2.5 * np.log10(sum(sum(image_HD))) + source_info['zero_point']
    magnification = source_info[
        'magnitude'] - mag_LS  #rough estimation in the magnitude change

    #DES - PSF convolution
    npsf = inter_psf(lens_info['psf'], lens_info['deltapix'],
                     source_info['deltapix'])
    source_conv = signal.fftconvolve(image_HD, npsf, mode='same')

    #resize in deltapix
    source_lensed_res = block_reduce(
        source_conv, source_info['HR_factor'])  #in case an HR_factor was used
    eq_pa = int(
        lens_info['deltapix'] * len(lens_info['image']) *
        len(source_lensed_res) /
        (source_info['deltapix'] * len(source_info['image']))
    )  #this value is to consider a source image of the same size in arcseconds as the lens
    bs = np.zeros([eq_pa, eq_pa])
    val = int((len(bs) - len(source_lensed_res)) / 2)
    bs[val:val + len(source_lensed_res),
       val:val + len(source_lensed_res)] = source_lensed_res
    source_size_scale = block_reduce(bs,
                                     int(len(bs) / len(lens_info['image'])))

    #flux rescale
    flux_img = sum(image_HD.flatten()) * (10**(
        0.4 * (lens_info['zero_point'] - source_info['zero_point'])))
    sc = sum(source_conv.flatten()) / flux_img
    source_scaled = source_size_scale / sc

    #cut right pixels size
    lens_size = min(np.shape(lens_info['image']))
    source_size = min(np.shape(source_scaled))

    if lens_size > source_size:
        xin = yin = int((lens_size - source_size) / 2)
        lens_final = lens_info['image'][xin:xin + source_size,
                                        yin:yin + source_size]
        source_final = source_scaled

    if lens_size < source_size:
        xin = yin = math.ceil((source_size - lens_size) / 2)
        source_final = source_scaled[xin:xin + lens_size, yin:yin + lens_size]
        lens_final = lens_info['image']
    else:
        lens_final = lens_info['image']
        source_final = source_scaled
    final = lens_final + source_final

    phot_ap = 2
    _, mag_sim = ap_phot(
        len(final) / 2,
        len(final) / 2, final, phot_ap / (lens_info['deltapix']),
        phot_ap / (lens_info['deltapix']), np.pi / 2., lens_info['zero_point'])

    #in the old code some images cause problems with some inf pixels... this will discard that system
    if magnification == np.float('-inf') or magnification == np.float('inf'):
        print('INFINITE MAGNIFICATION')
        magnification = 10000

    return {
        'simulation': final,
        'src_image': source_info['image'],
        'mag_sim': mag_sim,
        'mag_lensed_src': mag_LS,
        'image_HD': image_HD,
        'resize': source_size_scale,
        'conv': source_conv,
        'magnification': magnification
    }
Esempio n. 6
0
    def sim_image(self, info_dict):
        """
        Simulate an image based on specifications in sim_dict
        
        :param info_dict: A single element of the output form Organizer.breakup()
        """
        output_image = []
        if self.return_planes:
            output_source, output_lens, output_point_source, output_noise = [], [], [], []
        output_metadata = []
        
        #set the cosmology
        cosmology_info = ['H0', 'Om0', 'Tcmb0', 'Neff', 'm_nu', 'Ob0']
        cosmo = FlatLambdaCDM(**dict_select_choose(list(info_dict.values())[0], cosmology_info))
        
        for band, sim_dict in info_dict.items():
            ### Geometry-independent image properties
            
            # Single Band Info
            single_band_info = ['read_noise', 'pixel_scale', 'ccd_gain', 'exposure_time',
                                'sky_brightness', 'seeing', 'magnitude_zero_point',
                                'num_exposures', 'data_count_unit', 'background_noise']
            kwargs_single_band = dict_select_choose(sim_dict, single_band_info)
            
            # Extinction
            extinction_class = None  #figure this out later
            
            # Numerics -- only use single number kwargs
            numerics_info = ['supersampling_factor', 'compute_mode', 'supersampling_convolution',
                             'supersampling_kernel_size', 'point_source_supersampling_factor', 
                             'convolution_kernel_size']
            kwargs_numerics = dict_select_choose(sim_dict, numerics_info)
            
            ### Geometry-dependent image properties
            
            # Dictionary for physical model
            kwargs_model = {'lens_model_list': [],  # list of lens models to be used
                            'lens_redshift_list': [],  # list of redshift of the deflections
                            'lens_light_model_list': [],  # list of unlensed light models to be used
                            'source_light_model_list': [],  # list of extended source models to be used
                            'source_redshift_list': [],  # list of redshfits of the sources in same order as source_light_model_list
                            'point_source_model_list': [], # list of point source models
                            'cosmo': cosmo,
                            'z_source': 0.0}

            # Lists for model kwargs
            kwargs_source_list, kwargs_lens_model_list, kwargs_lens_light_list, kwargs_point_source_list = [], [], [], []
            
            for plane_num in range(1, sim_dict['NUMBER_OF_PLANES'] + 1):
                
                for obj_num in range(1, sim_dict['PLANE_{0}-NUMBER_OF_OBJECTS'.format(plane_num)] + 1):
                    
                    prefix = 'PLANE_{0}-OBJECT_{1}-'.format(plane_num, obj_num)
                    
                    ### Point sources
                    if sim_dict[prefix + 'HOST'] != 'None':
                        
                        # Foreground objects
                        if sim_dict[prefix + 'HOST'] == 'Foreground':
                            kwargs_model['point_source_model_list'].append('UNLENSED')
                            ps_info = ['PLANE_{0}-OBJECT_{1}-{2}'.format(plane_num, obj_num, x) for x in ['ra_image', 'dec_image', 'magnitude']]
                            ps_dict_info = dict_select(sim_dict, ps_info)
                            kwargs_point_source_list.append({'ra_image': [ps_dict_info[prefix + 'ra_image']], 'dec_image': [ps_dict_info[prefix + 'dec_image']], 'magnitude': [ps_dict_info[prefix + 'magnitude']]})
                        # Real point sources
                        else:
                            if plane_num < sim_dict['NUMBER_OF_PLANES']:
                                # point sources in the lens
                                kwargs_model['point_source_model_list'].append('LENSED_POSITION')
                                ps_info = ['PLANE_{0}-OBJECT_{1}-{2}'.format(plane_num, obj_num, x) for x in ['ra', 'dec', 'magnitude']]
                                ps_dict_info = dict_select(sim_dict, ps_info)
                                kwargs_point_source_list.append({'ra_image': [ps_dict_info[prefix + 'ra']],
                                                                 'dec_image': [ps_dict_info[prefix + 'dec']],
                                                                 'magnitude': [ps_dict_info[prefix + 'magnitude']]})
                            elif plane_num == sim_dict['NUMBER_OF_PLANES']:
                                # point sources in the source plane
                                kwargs_model['point_source_model_list'].append('SOURCE_POSITION')
                                ps_info = ['PLANE_{0}-OBJECT_{1}-{2}'.format(plane_num, obj_num, x) for x in ['ra', 'dec', 'magnitude']]
                                ps_dict_info = dict_select(sim_dict, ps_info)
                                kwargs_point_source_list.append({'ra_source': ps_dict_info['PLANE_{0}-OBJECT_{1}-ra'.format(plane_num, obj_num)], 
                                                                 'dec_source': ps_dict_info['PLANE_{0}-OBJECT_{1}-dec'.format(plane_num, obj_num)], 
                                                                 'magnitude': ps_dict_info['PLANE_{0}-OBJECT_{1}-magnitude'.format(plane_num, obj_num)]})
                            else:
                                #should never get here
                                assert False
                                
                        # the number of profiles will be zero for point sources, so just skip ahead
                        continue
                    
                    ### Model keywords
                    # All planes except last one - treat as lens
                    if plane_num < sim_dict['NUMBER_OF_PLANES']:
                        for light_profile_num in range(1, sim_dict[prefix + 'NUMBER_OF_LIGHT_PROFILES'] +1):
                            kwargs_model['lens_light_model_list'].append(sim_dict[prefix + 'LIGHT_PROFILE_{0}-NAME'.format(light_profile_num)])
                            #kwargs_model['lens_redshift_list'].append(sim_dict[prefix + 'REDSHIFT'])
                            kwargs_lens_light_list.append(select_params(sim_dict, prefix + 'LIGHT_PROFILE_{0}-'.format(light_profile_num)))
                        for mass_profile_num in range(1, sim_dict[prefix + 'NUMBER_OF_MASS_PROFILES'] +1):
                            kwargs_model['lens_model_list'].append(sim_dict[prefix + 'MASS_PROFILE_{0}-NAME'.format(mass_profile_num)])
                            kwargs_model['lens_redshift_list'].append(sim_dict[prefix + 'REDSHIFT'])
                            mass_params = select_params(sim_dict, prefix + 'MASS_PROFILE_{0}-'.format(mass_profile_num))
                            kwargs_lens_model_list.append(mass_params)
                            if 'sigma_v' in mass_params.keys(): # save simga_v locations so that we can calculate theta_E for the metadata
                                output_metadata.append({'PARAM_NAME':  prefix + 'MASS_PROFILE_{0}-sigma_v-{1}'.format(mass_profile_num, band),
                                                        'PARAM_VALUE': mass_params['sigma_v'],
                                                        'LENS_MODEL_IDX': len(kwargs_lens_model_list) - 1})
                        for shear_profile_num in range(1, sim_dict[prefix + 'NUMBER_OF_SHEAR_PROFILES'] +1):
                            kwargs_model['lens_model_list'].append(sim_dict[prefix + 'SHEAR_PROFILE_{0}-NAME'.format(shear_profile_num)])
                            kwargs_model['lens_redshift_list'].append(sim_dict[prefix + 'REDSHIFT'])
                            mass_params = select_params(sim_dict, prefix + 'SHEAR_PROFILE_{0}-'.format(shear_profile_num))
                            kwargs_lens_model_list.append(select_params(sim_dict, prefix + 'SHEAR_PROFILE_{0}-'.format(shear_profile_num)))
                                                   
                    # Last Plane - treat as source
                    elif plane_num == sim_dict['NUMBER_OF_PLANES']:
                        kwargs_model['z_source'] = sim_dict[prefix + 'REDSHIFT']
                        for light_profile_num in range(1, sim_dict[prefix + 'NUMBER_OF_LIGHT_PROFILES'] +1):
                            kwargs_model['source_light_model_list'].append(sim_dict[prefix + 'LIGHT_PROFILE_{0}-NAME'.format(light_profile_num)])
                            kwargs_model['source_redshift_list'].append(sim_dict[prefix + 'REDSHIFT'])
                            kwargs_source_list.append(select_params(sim_dict, prefix + 'LIGHT_PROFILE_{0}-'.format(light_profile_num)))
                        
                    else:
                        # Should never get here
                        assert False

            # Make image
            sim = SimAPI(numpix=sim_dict['numPix'], 
                         kwargs_single_band=kwargs_single_band, 
                         kwargs_model=kwargs_model)

            imSim = sim.image_model_class(kwargs_numerics)
            
            kwargs_lens_light_list, kwargs_source_list, kwargs_point_source_list = sim.magnitude2amplitude(kwargs_lens_light_mag=kwargs_lens_light_list,
                                                                                                           kwargs_source_mag=kwargs_source_list,
                                                                                                           kwargs_ps_mag=kwargs_point_source_list)
            
            kwargs_lens_model_list = sim.physical2lensing_conversion(kwargs_mass=kwargs_lens_model_list)
            # Save theta_E (and sigma_v if used)
            for ii in range(len(output_metadata)):
                output_metadata.append({'PARAM_NAME': output_metadata[ii]['PARAM_NAME'].replace('sigma_v', 'theta_E'),
                                        'PARAM_VALUE': kwargs_lens_model_list[output_metadata[ii]['LENS_MODEL_IDX']]['theta_E'],
                                        'LENS_MODEL_IDX': output_metadata[ii]['LENS_MODEL_IDX']})
                
            try:
                image = imSim.image(kwargs_lens=kwargs_lens_model_list,
                                    kwargs_lens_light=kwargs_lens_light_list,
                                    kwargs_source=kwargs_source_list,
                                    kwargs_ps=kwargs_point_source_list)
            except Exception:
                # Some sort of lenstronomy error
                print("kwargs_numerics", kwargs_numerics)
                print("kwargs_single_band", kwargs_single_band)
                print("kwargs_model", kwargs_model)
                print("kwargs_lens_model_list", kwargs_lens_model_list)
                print("kwargs_lens_light_list", kwargs_lens_light_list)
                print("kwargs_source_list", kwargs_source_list)
                print("kwargs_point_source_list", kwargs_point_source_list)
                assert False
                                
            # Solve lens equation if desired
            if self.solve_lens_equation:
                solver = lens_equation_solver.LensEquationSolver(imSim.LensModel)
                x_mins, y_mins = solver.image_position_from_source(sourcePos_x=kwargs_source_list[0]['center_x'],
                                                                   sourcePos_y=kwargs_source_list[0]['center_y'],
                                                                   kwargs_lens=kwargs_lens_model_list)
                num_source_images = len(x_mins)
            
            # Add noise
            image_noise = np.zeros(np.shape(image))
            for noise_source_num in range(1, sim_dict['NUMBER_OF_NOISE_SOURCES'] + 1):
                image_noise += self.generate_noise(sim_dict['NOISE_SOURCE_{0}-NAME'.format(noise_source_num)],
                                                   np.shape(image),
                                                   select_params(sim_dict, 'NOISE_SOURCE_{0}-'.format(noise_source_num)))
            image += image_noise
                
            # Combine with other bands
            output_image.append(image)

            # Store plane-separated info if requested
            if self.return_planes:
                output_lens.append(imSim.lens_surface_brightness(kwargs_lens_light_list))
                output_source.append(imSim.source_surface_brightness(kwargs_source_list, kwargs_lens_model_list))
                output_point_source.append(imSim.point_source(kwargs_point_source_list, kwargs_lens_model_list))
                output_noise.append(image_noise)
        
        # Return the desired information in a dictionary
        return_dict = {'output_image': np.array(output_image),
                       'output_lens_plane': None,
                       'output_source_plane': None,
                       'output_point_source_plane': None,
                       'output_noise_plane': None,
                       'x_mins': None,
                       'y_mins': None,
                       'num_source_images': None,
                       'additional_metadata': output_metadata}
        if self.return_planes:
            return_dict['output_lens_plane'] = np.array(output_lens)
            return_dict['output_source_plane'] = np.array(output_source)
            return_dict['output_point_source_plane'] = np.array(output_point_source)
            return_dict['output_noise_plane'] = np.array(output_noise)
        if self.solve_lens_equation:
            return_dict['x_mins'] = x_mins
            return_dict['y_mins'] = y_mins
            return_dict['num_source_images'] = num_source_images

        return return_dict
def sim_single_image(sn_data, geo_dict, label):

    DECam = {
        'read_noise':
        10,  # std of noise generated by read-out (in units of electrons)                                                   
        'pixel_scale':
        0.263,  # scale (in arcseonds) of pixels                                                                               
        'ccd_gain':
        4.5  # electrons/ADU (analog-to-digital unit).                                                                      
    }

    obs_info = {
        'exposure_time': 90,
        'sky_brightness': sn_data['SKYMAG'],
        'magnitude_zero_point': sn_data['ZPTMAG'],
        'num_exposures': 1,
        'seeing': sn_data['PSF'] / 3,
        'psf_type': 'GAUSSIAN'
    }

    kwargs_model = {
        'lens_model_list': [
            'SIE', 'SHEAR'
        ],  # list of lens models to be used                                                         
        'lens_light_model_list': [
            'SERSIC_ELLIPSE'
        ],  # list of unlensed light models to be used                                           
        'source_light_model_list': [
            'SERSIC_ELLIPSE'
        ],  # list of extended source models to be used                                        
        'point_source_model_list': [
            'SOURCE_POSITION'
        ],  # list of point source models to be used                                          
    }

    numpix = 64
    kwargs_merged = util.merge_dicts(DECam, obs_info)
    kwargs_numerics = {'point_source_supersampling_factor': 1}

    sim = SimAPI(numpix=numpix,
                 kwargs_single_band=kwargs_merged,
                 kwargs_model=kwargs_model,
                 kwargs_numerics=kwargs_numerics)

    imSim = sim.image_model_class

    x_grid, y_grid = util.make_grid(numPix=10, deltapix=1)
    flux = gauss.function(x_grid,
                          y_grid,
                          amp=1,
                          sigma=2,
                          e1=0.4,
                          e2=0,
                          center_x=0,
                          center_y=0)
    image_gauss = util.array2image(flux)

    ## Set point source mag
    geo_dict['ps_kwargs']['magnitude'] = sn_data['ABMAG'] * 1.4

    if label == 1:
        geo_dict = force_larger_ps_separation(geo_dict)
    elif label == 2:
        geo_dict = force_no_lensing_with_ps(geo_dict)
    elif label == 3:
        geo_dict = force_lensed_agn(geo_dict)
    elif label == 4:
        geo_dict = force_non_lensed_agn(geo_dict)
    else:
        print("Unexpected option passed as class label")
        sys.exit()

    #lens light
    kwargs_lens_light_mag = [geo_dict['lens_kwargs']]
    # source light
    kwargs_source_mag = [geo_dict['source_kwargs']]
    # point source
    kwargs_ps_mag = [geo_dict['ps_kwargs']]

    kwargs_lens_light, kwargs_source, kwargs_ps = sim.magnitude2amplitude(
        kwargs_lens_light_mag, kwargs_source_mag, kwargs_ps_mag)
    kwargs_lens = [
        geo_dict['lens_model_kwargs_sie'], geo_dict['lens_model_kwargs_shear']
    ]

    image = imSim.image(kwargs_lens, kwargs_source, kwargs_lens_light,
                        kwargs_ps)
    image += sim.noise_for_model(model=image)

    return np.log10(image + 2)
class PointSourceVariability(object):
    """
    This class enables to plug in a variable point source in the source plane to be added on top of a fixed lens and
    extended surface brightness model. The class inherites SimAPI and additionally requires the lens and light model
    parameters as well as a position in the source plane.

    The intrinsic source variability can be defined by the user and additional uncorrelated variability in the image
    plane can be plugged in as well (e.g. due to micro-lensing)
    """
    def __init__(self,
                 source_x,
                 source_y,
                 variability_func,
                 numpix,
                 kwargs_single_band,
                 kwargs_model,
                 kwargs_numerics,
                 kwargs_lens,
                 kwargs_source_mag=None,
                 kwargs_lens_light_mag=None,
                 kwargs_ps_mag=None):
        """

        :param source_x:
        :param source_y:
        :param variability_func: function that returns a brightness (in magnitude) as a function of time t
        :param numpix:
        :param kwargs_single_band:
        :param kwargs_model:
        :param kwargs_numerics:
        :param kwargs_lens:
        :param kwargs_source_mag:
        :param kwargs_lens_light_mag:
        :param kwargs_ps_mag:
        """

        # create background SimAPI class instance
        self.sim_api_bkg = SimAPI(numpix, kwargs_single_band, kwargs_model,
                                  kwargs_numerics)
        kwargs_lens_light, kwargs_source, kwargs_ps = self.sim_api_bkg.magnitude2amplitude(
            kwargs_lens_light_mag, kwargs_source_mag, kwargs_ps_mag)
        self._image_bkg = self.sim_api_bkg.image_model_class.image(
            kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps)
        # compute image positions of point source
        x_center, y_center = self.sim_api_bkg.data_class.center
        search_window = self.sim_api_bkg.data_class.width
        lensModel = self.sim_api_bkg.image_model_class.LensModel
        solver = LensEquationSolver(lensModel=lensModel)
        image_x, image_y = solver.image_position_from_source(
            source_x,
            source_y,
            kwargs_lens,
            min_distance=0.1,
            search_window=search_window,
            precision_limit=10**(-10),
            num_iter_max=100,
            arrival_time_sort=True,
            x_center=x_center,
            y_center=y_center)
        mag = lensModel.magnification(image_x, image_y, kwargs_lens)
        dt_days = lensModel.arrival_time(image_x, image_y, kwargs_lens)
        dt_days -= np.min(
            dt_days
        )  # shift the arrival times such that the first image arrives at t=0 and the other
        #  times at t>=0
        # add image plane source model
        kwargs_model_ps = {'point_source_model_list': ['LENSED_POSITION']}
        self.sim_api_ps = SimAPI(numpix, kwargs_single_band, kwargs_model_ps,
                                 kwargs_numerics)
        self._kwargs_lens = kwargs_lens
        self._dt_days = dt_days
        self._mag = mag
        self._image_x, self._image_y = image_x, image_y
        self._variability_func = variability_func

    @property
    def delays(self):
        """

        :return: time delays
        """
        return self._dt_days

    @property
    def image_bkg(self):
        """

        :return: 2d numpy array, image of the extended light components without the variable source
        """
        return self._image_bkg

    def image_time(self, time=0):
        """

        :param time: time relative to the definition of t=0 for the first appearing image
        :return: image with time variable source at given time
        """
        kwargs_ps_time = self.point_source_time(time)
        point_source = self.sim_api_ps.image_model_class.point_source(
            kwargs_ps_time, kwargs_lens=self._kwargs_lens)
        return point_source + self.image_bkg

    def point_source_time(self, t):
        """

        :param t: time (in units of days)
        :return: image plane parameters of the point source observed at t
        """
        mag = np.zeros_like(self._dt_days)
        kwargs_ps = [{'ra_image': self._image_x, 'dec_image': self._image_y}]
        for i, dt in enumerate(self._dt_days):
            t_i = -dt + t
            mag[i] = self._variability_func(t_i)
        kwargs_ps[0]['point_amp'] = self.sim_api_ps.magnitude2cps(
            mag) * np.abs(self._mag)
        return kwargs_ps
    def sim_image(self, info_dict):
        """
        Simulate an image based on specifications in sim_dict
        
        Args:
            info_dict (dict): A single element from the list produced interanlly by input_reader.Organizer.breakup(). 
                Contains all the properties of a single image to generate.
        """
        output_image = []
        if self.return_planes:
            output_source, output_lens, output_point_source, output_noise = [], [], [], []
        output_metadata = []

        #set the cosmology
        cosmology_info = ['H0', 'Om0', 'Tcmb0', 'Neff', 'm_nu', 'Ob0']
        cosmo = FlatLambdaCDM(
            **dict_select_choose(list(info_dict.values())[0], cosmology_info))

        for band, sim_dict in info_dict.items():

            # Parse the info dict
            params = self.parse_single_band_info_dict(sim_dict,
                                                      cosmo,
                                                      band=band)
            kwargs_single_band = params[0]
            kwargs_model = params[1]
            kwargs_numerics = params[2]
            kwargs_lens_light_list = params[3]
            kwargs_source_list = params[4]
            kwargs_point_source_list = params[5]
            kwargs_lens_model_list = params[6]
            output_metadata += params[7]

            # Make image
            # data properties
            kwargs_data = sim_util.data_configure_simple(
                sim_dict['numPix'], kwargs_single_band['pixel_scale'],
                kwargs_single_band['exposure_time'])
            data_class = ImageData(**kwargs_data)

            # psf properties
            kwargs_psf = {
                'psf_type': kwargs_single_band['psf_type'],
                'pixel_size': kwargs_single_band['pixel_scale'],
                'fwhm': kwargs_single_band['seeing']
            }
            psf_class = PSF(**kwargs_psf)

            # SimAPI instance for conversion to observed quantities
            sim = SimAPI(numpix=sim_dict['numPix'],
                         kwargs_single_band=kwargs_single_band,
                         kwargs_model=kwargs_model)
            kwargs_lens_model_list = sim.physical2lensing_conversion(
                kwargs_mass=kwargs_lens_model_list)
            kwargs_lens_light_list, kwargs_source_list, _ = sim.magnitude2amplitude(
                kwargs_lens_light_mag=kwargs_lens_light_list,
                kwargs_source_mag=kwargs_source_list)

            # lens model properties
            lens_model_class = LensModel(
                lens_model_list=kwargs_model['lens_model_list'],
                z_lens=kwargs_model['lens_redshift_list'][0],
                z_source=kwargs_model['z_source'],
                cosmo=cosmo)

            # source properties
            source_model_class = LightModel(
                light_model_list=kwargs_model['source_light_model_list'])

            # lens light properties
            lens_light_model_class = LightModel(
                light_model_list=kwargs_model['lens_light_model_list'])

            # solve for PS positions to incorporate time delays
            lensEquationSolver = LensEquationSolver(lens_model_class)
            kwargs_ps = []
            for ps_idx, ps_mag in enumerate(kwargs_point_source_list):

                # modify the SimAPI instance to do one point source at a time
                temp_kwargs_model = {k: v for k, v in kwargs_model.items()}
                temp_kwargs_model['point_source_model_list'] = [
                    kwargs_model['point_source_model_list'][ps_idx]
                ]
                sim = SimAPI(numpix=sim_dict['numPix'],
                             kwargs_single_band=kwargs_single_band,
                             kwargs_model=temp_kwargs_model)

                if kwargs_model['point_source_model_list'][
                        ps_idx] == 'SOURCE_POSITION':
                    # convert each image to an amplitude
                    amplitudes = []
                    for mag in ps_mag['magnitude']:
                        ps_dict = {k: v for k, v in ps_mag.items()}
                        ps_dict['magnitude'] = mag
                        _, _2, ps = sim.magnitude2amplitude(
                            kwargs_ps_mag=[ps_dict])
                        amplitudes.append(ps[0]['source_amp'])

                    x_image, y_image = lensEquationSolver.findBrightImage(
                        ps[0]['ra_source'],
                        ps[0]['dec_source'],
                        kwargs_lens_model_list,
                        numImages=4,  # max number of images
                        min_distance=kwargs_single_band['pixel_scale'],
                        search_window=sim_dict['numPix'] *
                        kwargs_single_band['pixel_scale'])
                    magnification = lens_model_class.magnification(
                        x_image, y_image, kwargs=kwargs_lens_model_list)
                    #amplitudes = np.array(amplitudes) * np.abs(magnification)
                    amplitudes = np.array(
                        [a * m for a, m in zip(amplitudes, magnification)])

                    kwargs_ps.append({
                        'ra_image': x_image,
                        'dec_image': y_image,
                        'point_amp': amplitudes
                    })

                else:
                    _, _2, ps = sim.magnitude2amplitude(kwargs_ps_mag=[ps_mag])
                    kwargs_ps.append(ps[0])

            # point source properties
            point_source_class = PointSource(point_source_type_list=[
                x if x != 'SOURCE_POSITION' else 'LENSED_POSITION'
                for x in kwargs_model['point_source_model_list']
            ],
                                             fixed_magnification_list=[False] *
                                             len(kwargs_ps))

            # create an image model
            image_model = ImageModel(data_class,
                                     psf_class,
                                     lens_model_class,
                                     source_model_class,
                                     lens_light_model_class,
                                     point_source_class,
                                     kwargs_numerics=kwargs_numerics)

            # generate image
            image_sim = image_model.image(kwargs_lens_model_list,
                                          kwargs_source_list,
                                          kwargs_lens_light_list, kwargs_ps)
            poisson = image_util.add_poisson(
                image_sim, exp_time=kwargs_single_band['exposure_time'])
            sigma_bkg = data_util.bkg_noise(
                kwargs_single_band['read_noise'],
                kwargs_single_band['exposure_time'],
                kwargs_single_band['sky_brightness'],
                kwargs_single_band['pixel_scale'],
                num_exposures=kwargs_single_band['num_exposures'])
            bkg = image_util.add_background(image_sim, sigma_bkd=sigma_bkg)
            image = image_sim + bkg + poisson

            # Save theta_E (and sigma_v if used)
            for ii in range(len(output_metadata)):
                output_metadata.append({
                    'PARAM_NAME':
                    output_metadata[ii]['PARAM_NAME'].replace(
                        'sigma_v', 'theta_E'),
                    'PARAM_VALUE':
                    kwargs_lens_model_list[output_metadata[ii]
                                           ['LENS_MODEL_IDX']]['theta_E'],
                    'LENS_MODEL_IDX':
                    output_metadata[ii]['LENS_MODEL_IDX']
                })

            # Solve lens equation if desired
            if self.solve_lens_equation:
                #solver = lens_equation_solver.LensEquationSolver(imSim.LensModel)
                #x_mins, y_mins = solver.image_position_from_source(sourcePos_x=kwargs_source_list[0]['center_x'],
                #                                                   sourcePos_y=kwargs_source_list[0]['center_y'],
                #                                                   kwargs_lens=kwargs_lens_model_list)
                x_mins, y_mins = x_image, y_image
                num_source_images = len(x_mins)

            # Add noise
            image_noise = np.zeros(np.shape(image))
            for noise_source_num in range(
                    1, sim_dict['NUMBER_OF_NOISE_SOURCES'] + 1):
                image_noise += self._generate_noise(
                    sim_dict['NOISE_SOURCE_{0}-NAME'.format(noise_source_num)],
                    np.shape(image),
                    select_params(
                        sim_dict,
                        'NOISE_SOURCE_{0}-'.format(noise_source_num)))
            image += image_noise

            # Combine with other bands
            output_image.append(image)

            # Store plane-separated info if requested
            if self.return_planes:
                output_lens.append(
                    image_model.lens_surface_brightness(
                        kwargs_lens_light_list))
                output_source.append(
                    image_model.source_surface_brightness(
                        kwargs_source_list, kwargs_lens_model_list))
                output_point_source.append(
                    image_model.point_source(kwargs_ps,
                                             kwargs_lens_model_list))
                output_noise.append(image_noise)

        # Return the desired information in a dictionary
        return_dict = {
            'output_image': np.array(output_image),
            'output_lens_plane': None,
            'output_source_plane': None,
            'output_point_source_plane': None,
            'output_noise_plane': None,
            'x_mins': None,
            'y_mins': None,
            'num_source_images': None,
            'additional_metadata': output_metadata
        }
        if self.return_planes:
            return_dict['output_lens_plane'] = np.array(output_lens)
            return_dict['output_source_plane'] = np.array(output_source)
            return_dict['output_point_source_plane'] = np.array(
                output_point_source)
            return_dict['output_noise_plane'] = np.array(output_noise)
        if self.solve_lens_equation:
            return_dict['x_mins'] = x_mins
            return_dict['y_mins'] = y_mins
            return_dict['num_source_images'] = num_source_images

        return return_dict