def setup(self): np.random.seed(42) # data specifics sigma_bkg = 0.05 # background noise per pixel exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 50 # cutout pixel size deltaPix = 0.1 # pixel size in arcsec (area per pixel = deltaPix**2) fwhm = 0.5 # full width half max of PSF kwargs_model = {'lens_model_list': ['SPEP'], 'lens_light_model_list': ['SERSIC'], 'source_light_model_list': ['SERSIC_ELLIPSE'], 'point_source_model_list': ['SOURCE_POSITION'], 'fixed_magnification_list': [True]} # PSF specification kwargs_band = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg) data_class = ImageData(**kwargs_band) kwargs_psf = {'psf_type': 'GAUSSIAN', 'fwhm': fwhm, 'pixel_size': deltaPix} psf_class = PSF(**kwargs_psf) print(np.shape(psf_class.kernel_point_source), 'test kernel shape -') kwargs_spemd = {'theta_E': 1., 'gamma': 1.95, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1} self.kwargs_lens = [kwargs_spemd] kwargs_sersic = {'amp': 1/0.05**2., 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0} # 'SERSIC_ELLIPSE': elliptical Sersic profile kwargs_sersic_ellipse = {'amp': 1., 'R_sersic': .6, 'n_sersic': 3, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1} self.kwargs_lens_light = [kwargs_sersic] self.kwargs_source = [kwargs_sersic_ellipse] self.kwargs_ps = [{'ra_source': 0.55, 'dec_source': 0.02, 'source_amp': 1.}] # quasar point source position in the source plane and intrinsic brightness self.kwargs_cosmo = {'D_dt': 1000} kwargs_numerics = {'supersampling_factor': 1, 'supersampling_convolution': False, 'compute_mode': 'gaussian'} lens_model_class, source_model_class, lens_light_model_class, point_source_class, extinction_class = class_creator.create_class_instances(**kwargs_model) imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, extinction_class, kwargs_numerics=kwargs_numerics) image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps) data_class.update_data(image_sim) kwargs_band['image_data'] = image_sim self.data_class = data_class self.psf_class = psf_class self.kwargs_model = kwargs_model self.kwargs_numerics = { 'supersampling_factor': 1, 'supersampling_convolution': False} kwargs_constraints = { 'num_point_source_list': [4], 'solver_type': 'NONE', # 'PROFILE', 'PROFILE_SHEAR', 'ELLIPSE', 'CENTER' 'Ddt_sampling': True } def condition_definition(kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps=None, kwargs_special=None, kwargs_extinction=None): logL = 0 if kwargs_lens_light[0]['R_sersic'] > kwargs_source[0]['R_sersic']: logL -= 10**15 return logL kwargs_likelihood = {'force_no_add_image': True, 'source_marg': True, 'astrometric_likelihood': True, 'position_uncertainty': 0.004, 'check_solver': True, 'solver_tolerance': 0.001, 'check_positive_flux': True, 'flux_ratio_likelihood': True, 'prior_lens': [[0, 'theta_E', 1, 0.1]], 'condition_definition': condition_definition } self.kwargs_data = {'multi_band_list': [[kwargs_band, kwargs_psf, kwargs_numerics]], 'multi_band_type': 'single-band', 'time_delays_measured': np.ones(4), 'time_delays_uncertainties': np.ones(4), 'flux_ratios': np.ones(4), 'flux_ratio_errors': np.ones(4) } self.param_class = Param(self.kwargs_model, **kwargs_constraints) self.imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) self.Likelihood = LikelihoodModule(kwargs_data_joint=self.kwargs_data, kwargs_model=kwargs_model, param_class=self.param_class, **kwargs_likelihood)
def fit_qso_multiband(QSO_im_list, psf_ave_list, psf_std_list=None, source_params=None,ps_param=None, background_rms_list=[0.04]*5, pix_sz = 0.168, exp_time = 300., fix_n=None, image_plot = True, corner_plot=True, flux_ratio_plot=True, deep_seed = False, fixcenter = False, QSO_msk_list=None, QSO_std_list=None, tag = None, no_MCMC= False, pltshow = 1, new_band_seq=None): ''' A quick fit for the QSO image with (so far) single sersice + one PSF. The input psf noise is optional. Parameter -------- QSO_im: An array of the QSO image. psf_ave: The psf image. psf_std: The psf noise, optional. source_params: The prior for the source. Default is given. background_rms: default as 0.04 exp_time: default at 2400. deep_seed: if Ture, more mcmc steps will be performed. tag: The name tag for save the plot Return -------- Will output the fitted image (Set image_plot = True), the corner_plot and the flux_ratio_plot. source_result, ps_result, image_ps, image_host To do -------- ''' # data specifics need to set up based on the data situation background_rms_list = background_rms_list # background noise per pixel (Gaussian) exp_time = exp_time # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = len(QSO_im_list[0]) # cutout pixel size deltaPix = pix_sz psf_type = 'PIXEL' # 'gaussian', 'pixel', 'NONE' kernel_list = psf_ave_list if new_band_seq == None: new_band_seq= range(len(QSO_im_list)) # if psf_std_list is not None: # kwargs_numerics_list = [{'subgrid_res': 1, 'psf_subgrid': False, 'psf_error_map': True}] * len(QSO_im_list) #Turn on the PSF error map # else: kwargs_numerics_list = [{'supersampling_factor': 1, 'supersampling_convolution': False}] * len(QSO_im_list) if source_params is None: # here are the options for the host galaxy fitting fixed_source = [] kwargs_source_init = [] kwargs_source_sigma = [] kwargs_lower_source = [] kwargs_upper_source = [] # Disk component, as modelled by an elliptical Sersic profile if fix_n == None: fixed_source.append({}) # we fix the Sersic index to n=1 (exponential) kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': 2., 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.}) kwargs_source_sigma.append({'n_sersic': 0.5, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1}) kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': 0.3, 'center_x': -10, 'center_y': -10}) kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3., 'n_sersic': 7., 'center_x': 10, 'center_y': 10}) elif fix_n is not None: fixed_source.append({'n_sersic': fix_n}) kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': fix_n, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.}) kwargs_source_sigma.append({'n_sersic': 0.001, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1}) kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': fix_n, 'center_x': -10, 'center_y': -10}) kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3, 'n_sersic': fix_n, 'center_x': 10, 'center_y': 10}) source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source] else: source_params = source_params if ps_param is None: center_x = 0.0 center_y = 0.0 point_amp = QSO_im_list[0].sum()/2. fixed_ps = [{}] kwargs_ps = [{'ra_image': [center_x], 'dec_image': [center_y], 'point_amp': [point_amp]}] kwargs_ps_init = kwargs_ps kwargs_ps_sigma = [{'ra_image': [0.01], 'dec_image': [0.01]}] kwargs_lower_ps = [{'ra_image': [-10], 'dec_image': [-10]}] kwargs_upper_ps = [{'ra_image': [10], 'dec_image': [10]}] ps_param = [kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps] else: ps_param = ps_param kwargs_params = {'source_model': source_params, 'point_source_model': ps_param} #============================================================================== #Doing the QSO fitting #============================================================================== kwargs_data_list, data_class_list = [], [] for i in range(len(QSO_im_list)): kwargs_data_i = sim_util.data_configure_simple(numPix, deltaPix, exp_time, background_rms_list[i], inverse=True) kwargs_data_list.append(kwargs_data_i) data_class_list.append(ImageData(**kwargs_data_i)) kwargs_psf_list = [] psf_class_list = [] for i in range(len(QSO_im_list)): kwargs_psf_i = {'psf_type': psf_type, 'kernel_point_source': kernel_list[i]} kwargs_psf_list.append(kwargs_psf_i) psf_class_list.append(PSF(**kwargs_psf_i)) data_class_list[i].update_data(QSO_im_list[i]) light_model_list = ['SERSIC_ELLIPSE'] * len(source_params[0]) lightModel = LightModel(light_model_list=light_model_list) point_source_list = ['UNLENSED'] pointSource = PointSource(point_source_type_list=point_source_list) imageModel_list = [] for i in range(len(QSO_im_list)): kwargs_data_list[i]['image_data'] = QSO_im_list[i] # if QSO_msk_list is not None: # kwargs_numerics_list[i]['mask'] = QSO_msk_list[i] if QSO_std_list is not None: kwargs_data_list[i]['noise_map'] = QSO_std_list[i] # if psf_std_list is not None: # kwargs_psf_list[i]['psf_error_map'] = psf_std_list[i] image_band_list = [] for i in range(len(QSO_im_list)): imageModel_list.append(ImageModel(data_class_list[i], psf_class_list[i], source_model_class=lightModel, point_source_class=pointSource, kwargs_numerics=kwargs_numerics_list[i])) image_band_list.append([kwargs_data_list[i], kwargs_psf_list[i], kwargs_numerics_list[i]]) multi_band_list = [image_band_list[i] for i in range(len(QSO_im_list))] # numerical options and fitting sequences kwargs_model = { 'source_light_model_list': light_model_list, 'point_source_model_list': point_source_list } if fixcenter == False: kwargs_constraints = {'num_point_source_list': [1] } elif fixcenter == True: kwargs_constraints = {'joint_source_with_point_source': [[0, 0]], 'num_point_source_list': [1] } kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty" 'source_marg': False, #In likelihood_module.LikelihoodModule -- whether to fully invert the covariance matrix for marginalization 'check_positive_flux': True, 'image_likelihood_mask_list': [QSO_msk_list] } # mpi = False # MPI possible, but not supported through that notebook. # The Params for the fitting. kwargs_init: initial input. kwargs_sigma: The parameter uncertainty. kwargs_fixed: fixed parameters; #kwargs_lower,kwargs_upper: Lower and upper limits. kwargs_data_joint = {'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear'} # 'single-band', 'multi-linear', 'joint-linear' fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) if deep_seed == False: fitting_kwargs_list = [ ['PSO', {'sigma_scale': 0.8, 'n_particles': 80, 'n_iterations': 60, 'compute_bands': [True]+[False]*(len(QSO_im_list)-1)}], ['align_images', {'n_particles': 10, 'n_iterations': 10, 'compute_bands': [False]+[True]*(len(QSO_im_list)-1)}], ['PSO', {'sigma_scale': 0.8, 'n_particles': 100, 'n_iterations': 200, 'compute_bands': [True]*len(QSO_im_list)}], ['MCMC', {'n_burn': 10, 'n_run': 20, 'walkerRatio': 50, 'sigma_scale': .1}] ] elif deep_seed == True: fitting_kwargs_list = [ ['PSO', {'sigma_scale': 0.8, 'n_particles': 150, 'n_iterations': 60, 'compute_bands': [True]+[False]*(len(QSO_im_list)-1)}], ['align_images', {'n_particles': 20, 'n_iterations': 20, 'compute_bands': [False]+[True]*(len(QSO_im_list)-1)}], ['PSO', {'sigma_scale': 0.8, 'n_particles': 150, 'n_iterations': 200, 'compute_bands': [True]*len(QSO_im_list)}], ['MCMC', {'n_burn': 20, 'n_run': 40, 'walkerRatio': 50, 'sigma_scale': .1}] ] if no_MCMC == True: del fitting_kwargs_list[-1] start_time = time.time() # lens_result, source_result, lens_light_result, ps_result, cosmo_temp, chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list) chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list) lens_result, source_result, lens_light_result, ps_result, cosmo_temp = fitting_seq.best_fit() end_time = time.time() print(end_time - start_time, 'total time needed for computation') print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ') source_result_list, ps_result_list = [], [] image_reconstructed_list, error_map_list, image_ps_list, image_host_list, shift_RADEC_list=[], [], [], [],[] imageLinearFit_list = [] for k in range(len(QSO_im_list)): # this is the linear inversion. The kwargs will be updated afterwards imageLinearFit_k = ImageLinearFit(data_class_list[k], psf_class_list[k], source_model_class=lightModel, point_source_class=pointSource, kwargs_numerics=kwargs_numerics_list[k]) image_reconstructed_k, error_map_k, _, _ = imageLinearFit_k.image_linear_solve(kwargs_source=source_result, kwargs_ps=ps_result) imageLinearFit_list.append(imageLinearFit_k) [kwargs_data_k, kwargs_psf_k, kwargs_numerics_k] = fitting_seq.multi_band_list[k] # data_class_k = data_class_list[k] #ImageData(**kwargs_data_k) # psf_class_k = psf_class_list[k] #PSF(**kwargs_psf_k) # imageModel_k = ImageModel(data_class_k, psf_class_k, source_model_class=lightModel, # point_source_class=pointSource, kwargs_numerics=kwargs_numerics_list[k]) imageModel_k = imageModel_list[k] modelPlot = ModelPlot(multi_band_list[k], kwargs_model, lens_result, source_result, lens_light_result, ps_result, arrow_size=0.02, cmap_string="gist_heat", likelihood_mask=QSO_im_list[k]) print("source_result", 'for', "k", source_result) image_host_k = [] for i in range(len(source_result)): image_host_k.append(imageModel_list[k].source_surface_brightness(source_result,de_lensed=True,unconvolved=False, k=i)) image_ps_k = imageModel_k.point_source(ps_result) # let's plot the output of the PSO minimizer image_reconstructed_list.append(image_reconstructed_k) source_result_list.append(source_result) ps_result_list.append(ps_result) error_map_list.append(error_map_k) image_ps_list.append(image_ps_k) image_host_list.append(image_host_k) if 'ra_shift' in fitting_seq.multi_band_list[k][0].keys(): shift_RADEC_list.append([fitting_seq.multi_band_list[k][0]['ra_shift'], fitting_seq.multi_band_list[k][0]['dec_shift']]) else: shift_RADEC_list.append([0,0]) if image_plot: f, axes = plt.subplots(3, 3, figsize=(16, 16), sharex=False, sharey=False) modelPlot.data_plot(ax=axes[0,0], text="Data") modelPlot.model_plot(ax=axes[0,1]) modelPlot.normalized_residual_plot(ax=axes[0,2], v_min=-6, v_max=6) modelPlot.decomposition_plot(ax=axes[1,0], text='Host galaxy', source_add=True, unconvolved=True) modelPlot.decomposition_plot(ax=axes[1,1], text='Host galaxy convolved', source_add=True) modelPlot.decomposition_plot(ax=axes[1,2], text='All components convolved', source_add=True, lens_light_add=True, point_source_add=True) modelPlot.subtract_from_data_plot(ax=axes[2,0], text='Data - Point Source', point_source_add=True) modelPlot.subtract_from_data_plot(ax=axes[2,1], text='Data - host galaxy', source_add=True) modelPlot.subtract_from_data_plot(ax=axes[2,2], text='Data - host galaxy - Point Source', source_add=True, point_source_add=True) f.tight_layout() if tag is not None: f.savefig('{0}_fitted_image_band{1}.pdf'.format(tag,new_band_seq[k])) if pltshow == 0: plt.close() else: plt.show() if corner_plot==True and no_MCMC==False and k ==0: # here the (non-converged) MCMC chain of the non-linear parameters if not samples_mcmc == []: n, num_param = np.shape(samples_mcmc) plot = corner.corner(samples_mcmc, labels=param_mcmc, show_titles=True) if tag is not None: plot.savefig('{0}_para_corner.pdf'.format(tag)) if pltshow == 0: plt.close() else: plt.show() if flux_ratio_plot==True and no_MCMC==False: param = Param(kwargs_model, kwargs_fixed_source=source_params[2], kwargs_fixed_ps=fixed_ps, **kwargs_constraints) mcmc_new_list = [] labels_new = [r"Quasar flux", r"host_flux", r"source_x", r"source_y"] # transform the parameter position of the MCMC chain in a lenstronomy convention with keyword arguments # for i in range(len(samples_mcmc)/10): kwargs_lens_out, kwargs_light_source_out, kwargs_light_lens_out, kwargs_ps_out, kwargs_cosmo = param.getParams(samples_mcmc[i+ len(samples_mcmc)/10*9]) image_reconstructed, _, _, _ = imageLinearFit_list[k].image_linear_solve(kwargs_source=kwargs_light_source_out, kwargs_ps=kwargs_ps_out) image_ps = imageModel_list[k].point_source(kwargs_ps_out) flux_quasar = np.sum(image_ps) image_disk = imageModel_list[k].source_surface_brightness(kwargs_light_source_out,de_lensed=True,unconvolved=False, k=0) flux_disk = np.sum(image_disk) source_x = kwargs_ps_out[0]['ra_image'] source_y = kwargs_ps_out[0]['dec_image'] if flux_disk>0: mcmc_new_list.append([flux_quasar, flux_disk, source_x, source_y]) plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True) if tag is not None: plot.savefig('{0}_HOSTvsQSO_corner_band{1}.pdf'.format(tag,new_band_seq[k])) if pltshow == 0: plt.close() else: plt.show() errp_list = [] for k in range(len(QSO_im_list)): if QSO_std_list is None: errp_list.append(np.sqrt(data_class_list[k].C_D+np.abs(error_map_list[k]))) else: errp_list.append(np.sqrt(QSO_std_list[k]**2+np.abs(error_map_list[k]))) return source_result_list, ps_result_list, image_ps_list, image_host_list, errp_list, shift_RADEC_list, fitting_seq #fitting_seq.multi_band_list
cut = int((psf_data.shape[0] - psf_data.shape[1]) / 2) if cut > 0: psf_data = psf_data[cut:-cut, :] elif cut < 0: psf_data = psf_data[:, -cut:cut] psf_data /= psf_data.sum() #%% numPix = 89 light_model_list = ['SERSIC_ELLIPSE'] kwargs_psf_high_res = { 'psf_type': 'PIXEL', 'kernel_point_source': psf_data, 'pixel_size': deltaPix } kwargs_data_high_res = sim_util.data_configure_simple(numPix, deltaPix) data_class = ImageData(**kwargs_data_high_res) psf_class = PSF(**kwargs_psf_high_res) center_x, center_y = np.random.uniform( -0.2, 0.2) * deltaPix, np.random.uniform(-0.2, 0.2) * deltaPix point_amp = AGN_flux point_source_list = ['UNLENSED'] pointSource = PointSource(point_source_type_list=point_source_list) kwargs_ps = [{ 'ra_image': [center_x], 'dec_image': [center_y], 'point_amp': [point_amp] }] light_model_list = ['SERSIC_ELLIPSE'] lightModel = LightModel(light_model_list=light_model_list)
def fit_galaxy(galaxy_im, psf_ave, psf_std=None, source_params=None, background_rms=0.04, pix_sz = 0.08, exp_time = 300., fix_n=None, image_plot = True, corner_plot=True, deep_seed = False, galaxy_msk=None, galaxy_std=None, flux_corner_plot = False, tag = None, no_MCMC= False, pltshow = 1, return_Chisq = False, dump_result = False, pso_diag=False): ''' A quick fit for the QSO image with (so far) single sersice + one PSF. The input psf noise is optional. Parameter -------- galaxy_im: An array of the QSO image. psf_ave: The psf image. psf_std: The psf noise, optional. source_params: The prior for the source. Default is given. background_rms: default as 0.04 exp_time: default at 2400. deep_seed: if Ture, more mcmc steps will be performed. tag: The name tag for save the plot Return -------- Will output the fitted image (Set image_plot = True), the corner_plot and the flux_ratio_plot. source_result, ps_result, image_ps, image_host To do -------- ''' # data specifics need to set up based on the data situation background_rms = background_rms # background noise per pixel (Gaussian) exp_time = exp_time # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = len(galaxy_im) # cutout pixel size deltaPix = pix_sz if psf_ave is not None: psf_type = 'PIXEL' # 'gaussian', 'pixel', 'NONE' kernel = psf_ave # if psf_std is not None: # kwargs_numerics = {'subgrid_res': 1, 'psf_error_map': True} #Turn on the PSF error map # else: kwargs_numerics = {'supersampling_factor': 1, 'supersampling_convolution': False} if source_params is None: # here are the options for the host galaxy fitting fixed_source = [] kwargs_source_init = [] kwargs_source_sigma = [] kwargs_lower_source = [] kwargs_upper_source = [] # Disk component, as modelled by an elliptical Sersic profile if fix_n == None: fixed_source.append({}) # we fix the Sersic index to n=1 (exponential) kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': 2., 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.}) kwargs_source_sigma.append({'n_sersic': 0.5, 'R_sersic': 0.1, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1}) kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.01, 'n_sersic': 0.3, 'center_x': -10, 'center_y': -10}) kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3., 'n_sersic': 7., 'center_x': 10, 'center_y': 10}) elif fix_n is not None: fixed_source.append({'n_sersic': fix_n}) kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': fix_n, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.}) kwargs_source_sigma.append({'n_sersic': 0.001, 'R_sersic': 0.1, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1}) kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.01, 'n_sersic': fix_n, 'center_x': -10, 'center_y': -10}) kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3, 'n_sersic': fix_n, 'center_x': 10, 'center_y': 10}) source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source] else: source_params = source_params kwargs_params = {'source_model': source_params} #============================================================================== #Doing the QSO fitting #============================================================================== kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, background_rms, inverse=True) data_class = ImageData(**kwargs_data) if psf_ave is not None: kwargs_psf = {'psf_type': psf_type, 'kernel_point_source': kernel} else: kwargs_psf = {'psf_type': 'NONE'} psf_class = PSF(**kwargs_psf) data_class.update_data(galaxy_im) light_model_list = ['SERSIC_ELLIPSE'] * len(source_params[0]) lightModel = LightModel(light_model_list=light_model_list) kwargs_model = { 'source_light_model_list': light_model_list} # numerical options and fitting sequences kwargs_constraints = {} kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty" 'source_marg': False, #In likelihood_module.LikelihoodModule -- whether to fully invert the covariance matrix for marginalization 'check_positive_flux': True, 'image_likelihood_mask_list': [galaxy_msk] } kwargs_data['image_data'] = galaxy_im if galaxy_std is not None: kwargs_data['noise_map'] = galaxy_std if psf_std is not None: kwargs_psf['psf_error_map'] = psf_std image_band = [kwargs_data, kwargs_psf, kwargs_numerics] multi_band_list = [image_band] kwargs_data_joint = {'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear'} # 'single-band', 'multi-linear', 'joint-linear' fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) if deep_seed == False: fitting_kwargs_list = [ ['PSO', {'sigma_scale': 0.8, 'n_particles': 50, 'n_iterations': 50}], ['MCMC', {'n_burn': 10, 'n_run': 10, 'walkerRatio': 50, 'sigma_scale': .1}] ] elif deep_seed == True: fitting_kwargs_list = [ ['PSO', {'sigma_scale': 0.8, 'n_particles': 100, 'n_iterations': 80}], ['MCMC', {'n_burn': 10, 'n_run': 15, 'walkerRatio': 50, 'sigma_scale': .1}] ] elif deep_seed == 'very_deep': fitting_kwargs_list = [ ['PSO', {'sigma_scale': 0.8, 'n_particles': 150, 'n_iterations': 150}], ['MCMC', {'n_burn': 10, 'n_run': 20, 'walkerRatio': 50, 'sigma_scale': .1}] ] if no_MCMC == True: fitting_kwargs_list = [fitting_kwargs_list[0], ] start_time = time.time() chain_list = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_seq.best_fit() ps_result = kwargs_result['kwargs_ps'] source_result = kwargs_result['kwargs_source'] if no_MCMC == False: sampler_type, samples_mcmc, param_mcmc, dist_mcmc = chain_list[1] # chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list) # lens_result, source_result, lens_light_result, ps_result, cosmo_temp = fitting_seq.best_fit() end_time = time.time() print(end_time - start_time, 'total time needed for computation') print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ') # this is the linear inversion. The kwargs will be updated afterwards imageModel = ImageModel(data_class, psf_class, source_model_class=lightModel,kwargs_numerics=kwargs_numerics) imageLinearFit = ImageLinearFit(data_class=data_class, psf_class=psf_class, source_model_class=lightModel, kwargs_numerics=kwargs_numerics) image_reconstructed, error_map, _, _ = imageLinearFit.image_linear_solve(kwargs_source=source_result, kwargs_ps=ps_result) # image_host = [] #!!! The linear_solver before and after could have different result for very faint sources. # for i in range(len(source_result)): # image_host_i = imageModel.source_surface_brightness(source_result,de_lensed=True,unconvolved=False, k=i) # print("image_host_i", source_result[i]) # print("total flux", image_host_i.sum()) # image_host.append(image_host_i) # let's plot the output of the PSO minimizer modelPlot = ModelPlot(multi_band_list, kwargs_model, kwargs_result, arrow_size=0.02, cmap_string="gist_heat", likelihood_mask_list=[galaxy_msk]) if pso_diag == True: f, axes = chain_plot.plot_chain_list(chain_list,0) if pltshow == 0: plt.close() else: plt.show() reduced_Chisq = imageLinearFit.reduced_chi2(image_reconstructed, error_map) if image_plot: f, axes = plt.subplots(1, 3, figsize=(16, 16), sharex=False, sharey=False) modelPlot.data_plot(ax=axes[0]) modelPlot.model_plot(ax=axes[1]) modelPlot.normalized_residual_plot(ax=axes[2], v_min=-6, v_max=6) f.tight_layout() #f.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0., hspace=0.05) if tag is not None: f.savefig('{0}_fitted_image.pdf'.format(tag)) if pltshow == 0: plt.close() else: plt.show() image_host = [] for i in range(len(source_result)): image_host_i = imageModel.source_surface_brightness(source_result,de_lensed=True,unconvolved=False, k=i) # print("image_host_i", source_result[i]) # print("total flux", image_host_i.sum()) image_host.append(image_host_i) if corner_plot==True and no_MCMC==False: # here the (non-converged) MCMC chain of the non-linear parameters if not samples_mcmc == []: n, num_param = np.shape(samples_mcmc) plot = corner.corner(samples_mcmc, labels=param_mcmc, show_titles=True) if tag is not None: plot.savefig('{0}_para_corner.pdf'.format(tag)) if pltshow == 0: plt.close() else: plt.show() if flux_corner_plot ==True and no_MCMC==False: param = Param(kwargs_model, kwargs_fixed_source=source_params[2], **kwargs_constraints) mcmc_new_list = [] labels_new = ["host{0} flux".format(i) for i in range(len(source_params[0]))] for i in range(len(samples_mcmc)): kwargs_out = param.args2kwargs(samples_mcmc[i]) kwargs_light_source_out = kwargs_out['kwargs_source'] kwargs_ps_out = kwargs_out['kwargs_ps'] image_reconstructed, _, _, _ = imageLinearFit.image_linear_solve(kwargs_source=kwargs_light_source_out, kwargs_ps=kwargs_ps_out) fluxs = [] for j in range(len(source_params[0])): image_j = imageModel.source_surface_brightness(kwargs_light_source_out,unconvolved= False, k=j) fluxs.append(np.sum(image_j)) mcmc_new_list.append( fluxs ) if int(i/1000) > int((i-1)/1000) : print(len(samples_mcmc), "MCMC samplers in total, finished translate:", i ) plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True) if tag is not None: plot.savefig('{0}_HOSTvsQSO_corner.pdf'.format(tag)) if pltshow == 0: plt.close() else: plt.show() if galaxy_std is None: noise_map = np.sqrt(data_class.C_D+np.abs(error_map)) else: noise_map = np.sqrt(galaxy_std**2+np.abs(error_map)) if dump_result == True: if flux_corner_plot==True and no_MCMC==False: trans_paras = [source_params[2], mcmc_new_list, labels_new, 'source_params[2], mcmc_new_list, labels_new'] else: trans_paras = [] picklename= tag + '.pkl' best_fit = [source_result, image_host, 'source_result, image_host'] # pso_fit = [chain_list, param_list, 'chain_list, param_list'] # mcmc_fit = [samples_mcmc, param_mcmc, dist_mcmc, 'samples_mcmc, param_mcmc, dist_mcmc'] chain_list_result = [chain_list, 'chain_list'] pickle.dump([best_fit, chain_list_result, trans_paras], open(picklename, 'wb')) if return_Chisq == False: return source_result, image_host, noise_map elif return_Chisq == True: return source_result, image_host, noise_map, reduced_Chisq
def fit_qso(QSO_im, psf_ave, psf_std=None, source_params=None,ps_param=None, background_rms=0.04, pix_sz = 0.168, exp_time = 300., fix_n=None, image_plot = True, corner_plot=True, supersampling_factor = 2, flux_ratio_plot=False, deep_seed = False, fixcenter = False, QSO_msk=None, QSO_std=None, tag = None, no_MCMC= False, pltshow = 1, return_Chisq = False, dump_result = False, pso_diag=False): ''' A quick fit for the QSO image with (so far) single sersice + one PSF. The input psf noise is optional. Parameter -------- QSO_im: An array of the QSO image. psf_ave: The psf image. psf_std: The psf noise, optional. source_params: The prior for the source. Default is given. If [], means no Sersic light. background_rms: default as 0.04 exp_time: default at 2400. deep_seed: if Ture, more mcmc steps will be performed. tag: The name tag for save the plot Return -------- Will output the fitted image (Set image_plot = True), the corner_plot and the flux_ratio_plot. source_result, ps_result, image_ps, image_host To do -------- ''' # data specifics need to set up based on the data situation background_rms = background_rms # background noise per pixel (Gaussian) exp_time = exp_time # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = len(QSO_im) # cutout pixel size deltaPix = pix_sz psf_type = 'PIXEL' # 'gaussian', 'pixel', 'NONE' kernel = psf_ave kwargs_numerics = {'supersampling_factor': supersampling_factor, 'supersampling_convolution': False} if source_params is None: # here are the options for the host galaxy fitting fixed_source = [] kwargs_source_init = [] kwargs_source_sigma = [] kwargs_lower_source = [] kwargs_upper_source = [] if fix_n == None: fixed_source.append({}) # we fix the Sersic index to n=1 (exponential) kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': 2., 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.}) kwargs_source_sigma.append({'n_sersic': 0.5, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1}) kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': 0.3, 'center_x': -10, 'center_y': -10}) kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3., 'n_sersic': 7., 'center_x': 10, 'center_y': 10}) elif fix_n is not None: fixed_source.append({'n_sersic': fix_n}) kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': fix_n, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.}) kwargs_source_sigma.append({'n_sersic': 0.001, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1}) kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': fix_n, 'center_x': -10, 'center_y': -10}) kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3, 'n_sersic': fix_n, 'center_x': 10, 'center_y': 10}) source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source] else: source_params = source_params if ps_param is None: center_x = 0.0 center_y = 0.0 point_amp = QSO_im.sum()/2. fixed_ps = [{}] kwargs_ps = [{'ra_image': [center_x], 'dec_image': [center_y], 'point_amp': [point_amp]}] kwargs_ps_init = kwargs_ps kwargs_ps_sigma = [{'ra_image': [0.05], 'dec_image': [0.05]}] kwargs_lower_ps = [{'ra_image': [-0.6], 'dec_image': [-0.6]}] kwargs_upper_ps = [{'ra_image': [0.6], 'dec_image': [0.6]}] ps_param = [kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps] else: ps_param = ps_param #============================================================================== #Doing the QSO fitting #============================================================================== kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, background_rms, inverse=True) data_class = ImageData(**kwargs_data) kwargs_psf = {'psf_type': psf_type, 'kernel_point_source': kernel} psf_class = PSF(**kwargs_psf) data_class.update_data(QSO_im) point_source_list = ['UNLENSED'] * len(ps_param[0]) pointSource = PointSource(point_source_type_list=point_source_list) if fixcenter == False: kwargs_constraints = {'num_point_source_list': [1] * len(ps_param[0]) } elif fixcenter == True: kwargs_constraints = {'joint_source_with_point_source': [[i, i] for i in range(len(ps_param[0]))], 'num_point_source_list': [1] * len(ps_param[0]) } if source_params == []: #fitting image as Point source only. kwargs_params = {'point_source_model': ps_param} lightModel = None kwargs_model = {'point_source_model_list': point_source_list } imageModel = ImageModel(data_class, psf_class, point_source_class=pointSource, kwargs_numerics=kwargs_numerics) kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty" 'image_likelihood_mask_list': [QSO_msk] } elif source_params != []: kwargs_params = {'source_model': source_params, 'point_source_model': ps_param} light_model_list = ['SERSIC_ELLIPSE'] * len(source_params[0]) lightModel = LightModel(light_model_list=light_model_list) kwargs_model = { 'source_light_model_list': light_model_list, 'point_source_model_list': point_source_list } imageModel = ImageModel(data_class, psf_class, source_model_class=lightModel, point_source_class=pointSource, kwargs_numerics=kwargs_numerics) # numerical options and fitting sequences kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty" 'source_marg': False, #In likelihood_module.LikelihoodModule -- whether to fully invert the covariance matrix for marginalization 'check_positive_flux': True, 'image_likelihood_mask_list': [QSO_msk] } kwargs_data['image_data'] = QSO_im if QSO_std is not None: kwargs_data['noise_map'] = QSO_std if psf_std is not None: kwargs_psf['psf_error_map'] = psf_std image_band = [kwargs_data, kwargs_psf, kwargs_numerics] multi_band_list = [image_band] kwargs_data_joint = {'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear'} # 'single-band', 'multi-linear', 'joint-linear' fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) if deep_seed == False: fitting_kwargs_list = [ ['PSO', {'sigma_scale': 0.8, 'n_particles': 100, 'n_iterations': 60}], ['MCMC', {'n_burn': 10, 'n_run': 10, 'walkerRatio': 50, 'sigma_scale': .1}] ] elif deep_seed == True: fitting_kwargs_list = [ ['PSO', {'sigma_scale': 0.8, 'n_particles': 250, 'n_iterations': 250}], ['MCMC', {'n_burn': 100, 'n_run': 200, 'walkerRatio': 10, 'sigma_scale': .1}] ] if no_MCMC == True: fitting_kwargs_list = [fitting_kwargs_list[0], ] start_time = time.time() chain_list = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_seq.best_fit() ps_result = kwargs_result['kwargs_ps'] source_result = kwargs_result['kwargs_source'] if no_MCMC == False: sampler_type, samples_mcmc, param_mcmc, dist_mcmc = chain_list[1] end_time = time.time() print(end_time - start_time, 'total time needed for computation') print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ') imageLinearFit = ImageLinearFit(data_class=data_class, psf_class=psf_class, source_model_class=lightModel, point_source_class=pointSource, kwargs_numerics=kwargs_numerics) image_reconstructed, error_map, _, _ = imageLinearFit.image_linear_solve(kwargs_source=source_result, kwargs_ps=ps_result) # this is the linear inversion. The kwargs will be updated afterwards modelPlot = ModelPlot(multi_band_list, kwargs_model, kwargs_result, arrow_size=0.02, cmap_string="gist_heat", likelihood_mask_list=[QSO_msk]) image_host = [] #!!! The linear_solver before and after LensModelPlot could have different result for very faint sources. for i in range(len(source_result)): image_host.append(imageModel.source_surface_brightness(source_result, de_lensed=True,unconvolved=False,k=i)) image_ps = [] for i in range(len(ps_result)): image_ps.append(imageModel.point_source(ps_result, k = i)) if pso_diag == True: f, axes = chain_plot.plot_chain_list(chain_list,0) if pltshow == 0: plt.close() else: plt.show() # let's plot the output of the PSO minimizer reduced_Chisq = imageLinearFit.reduced_chi2(image_reconstructed, error_map) if image_plot: f, axes = plt.subplots(3, 3, figsize=(16, 16), sharex=False, sharey=False) modelPlot.data_plot(ax=axes[0,0], text="Data") modelPlot.model_plot(ax=axes[0,1]) modelPlot.normalized_residual_plot(ax=axes[0,2], v_min=-6, v_max=6) modelPlot.decomposition_plot(ax=axes[1,0], text='Host galaxy', source_add=True, unconvolved=True) modelPlot.decomposition_plot(ax=axes[1,1], text='Host galaxy convolved', source_add=True) modelPlot.decomposition_plot(ax=axes[1,2], text='All components convolved', source_add=True, lens_light_add=True, point_source_add=True) modelPlot.subtract_from_data_plot(ax=axes[2,0], text='Data - Point Source', point_source_add=True) modelPlot.subtract_from_data_plot(ax=axes[2,1], text='Data - host galaxy', source_add=True) modelPlot.subtract_from_data_plot(ax=axes[2,2], text='Data - host galaxy - Point Source', source_add=True, point_source_add=True) f.tight_layout() #f.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0., hspace=0.05) if tag is not None: f.savefig('{0}_fitted_image.pdf'.format(tag)) if pltshow == 0: plt.close() else: plt.show() if corner_plot==True and no_MCMC==False: # here the (non-converged) MCMC chain of the non-linear parameters if not samples_mcmc == []: n, num_param = np.shape(samples_mcmc) plot = corner.corner(samples_mcmc, labels=param_mcmc, show_titles=True) if tag is not None: plot.savefig('{0}_para_corner.pdf'.format(tag)) plt.close() # if pltshow == 0: # plt.close() # else: # plt.show() if flux_ratio_plot==True and no_MCMC==False: param = Param(kwargs_model, kwargs_fixed_source=source_params[2], kwargs_fixed_ps=ps_param[2], **kwargs_constraints) mcmc_new_list = [] if len(ps_param[2]) == 1: labels_new = ["Quasar flux"] + ["host{0} flux".format(i) for i in range(len(source_params[0]))] else: labels_new = ["Quasar{0} flux".format(i) for i in range(len(ps_param[2]))] + ["host{0} flux".format(i) for i in range(len(source_params[0]))] if len(samples_mcmc) > 10000: trans_steps = [len(samples_mcmc)-10000, len(samples_mcmc)] else: trans_steps = [0, len(samples_mcmc)] for i in range(trans_steps[0], trans_steps[1]): kwargs_out = param.args2kwargs(samples_mcmc[i]) kwargs_light_source_out = kwargs_out['kwargs_source'] kwargs_ps_out = kwargs_out['kwargs_ps'] image_reconstructed, _, _, _ = imageLinearFit.image_linear_solve(kwargs_source=kwargs_light_source_out, kwargs_ps=kwargs_ps_out) flux_quasar = [] if len(ps_param[0]) == 1: image_ps_j = imageModel.point_source(kwargs_ps_out) flux_quasar.append(np.sum(image_ps_j)) else: for j in range(len(ps_param[0])): image_ps_j = imageModel.point_source(kwargs_ps_out, k=j) flux_quasar.append(np.sum(image_ps_j)) fluxs = [] for j in range(len(source_params[0])): image_j = imageModel.source_surface_brightness(kwargs_light_source_out,unconvolved= False, k=j) fluxs.append(np.sum(image_j)) mcmc_new_list.append(flux_quasar + fluxs ) if int(i/1000) > int((i-1)/1000) : print(len(samples_mcmc), "MCMC samplers in total, finished translate:", i ) plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True) if tag is not None: plot.savefig('{0}_HOSTvsQSO_corner.pdf'.format(tag)) if pltshow == 0: plt.close() else: plt.show() if QSO_std is None: noise_map = np.sqrt(data_class.C_D+np.abs(error_map)) else: noise_map = np.sqrt(QSO_std**2+np.abs(error_map)) if dump_result == True: if flux_ratio_plot==True and no_MCMC==False: trans_paras = [mcmc_new_list, labels_new, 'mcmc_new_list, labels_new'] else: trans_paras = [] picklename= tag + '.pkl' best_fit = [source_result, image_host, ps_result, image_ps,'source_result, image_host, ps_result, image_ps'] chain_list_result = [chain_list, 'chain_list'] kwargs_fixed_source=source_params[2] kwargs_fixed_ps=ps_param[2] classes = data_class, psf_class, lightModel, pointSource material = multi_band_list, kwargs_model, kwargs_result, QSO_msk, kwargs_fixed_source, kwargs_fixed_ps, kwargs_constraints, kwargs_numerics, classes pickle.dump([best_fit, chain_list_result, trans_paras, material], open(picklename, 'wb')) if return_Chisq == False: return source_result, ps_result, image_ps, image_host, noise_map elif return_Chisq == True: return source_result, ps_result, image_ps, image_host, noise_map, reduced_Chisq
def setup(self): # data specifics sigma_bkg = 0.05 # background noise per pixel exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 10 # cutout pixel size deltaPix = 0.1 # pixel size in arcsec (area per pixel = deltaPix**2) fwhm = 0.5 # full width half max of PSF # PSF specification kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg) data_class = ImageData(**kwargs_data) kwargs_psf_gaussian = { 'psf_type': 'GAUSSIAN', 'fwhm': fwhm, 'pixel_size': deltaPix } psf = PSF(**kwargs_psf_gaussian) kwargs_psf = { 'psf_type': 'PIXEL', 'kernel_point_source': psf.kernel_point_source } psf_class = PSF(**kwargs_psf) kwargs_spemd = { 'theta_E': 1., 'gamma': 1.8, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1 } lens_model_list = ['SPEP'] self.kwargs_lens = [kwargs_spemd] lens_model_class = LensModel(lens_model_list=lens_model_list) kwargs_sersic = { 'amp': 1., 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0 } # 'SERSIC_ELLIPSE': elliptical Sersic profile kwargs_sersic_ellipse = { 'amp': 1., 'R_sersic': .6, 'n_sersic': 3, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1 } lens_light_model_list = ['SERSIC'] self.kwargs_lens_light = [kwargs_sersic] lens_light_model_class = LightModel( light_model_list=lens_light_model_list) source_model_list = ['SERSIC_ELLIPSE'] self.kwargs_source = [kwargs_sersic_ellipse] source_model_class = LightModel(light_model_list=source_model_list) kwargs_numerics = { 'supersampling_factor': 1, 'supersampling_convolution': False, 'compute_mode': 'regular' } imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, kwargs_numerics=kwargs_numerics) image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light) data_class.update_data(image_sim) kwargs_data['image_data'] = image_sim kwargs_data_joint = { 'multi_band_list': [[kwargs_data, kwargs_psf, kwargs_numerics]], 'multi_band_type': 'single-band' } self.data_class = data_class self.psf_class = psf_class kwargs_model = { 'lens_model_list': lens_model_list, 'source_light_model_list': source_model_list, 'lens_light_model_list': lens_light_model_list, 'fixed_magnification_list': [False], } self.kwargs_numerics = {'subgrid_res': 1, 'psf_subgrid': False} kwargs_constraints = { 'image_plane_source_list': [False] * len(source_model_list) } kwargs_likelihood = { 'source_marg': False, 'image_position_uncertainty': 0.004, 'check_matched_source_position': False, 'source_position_tolerance': 0.001, 'source_position_sigma': 0.001, } # reduce number of param to sample (for runtime) kwargs_fixed_lens = [{ 'gamma': 1.8, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1 }] kwargs_lower_lens = [{'theta_E': 0.8}] kwargs_upper_lens = [{'theta_E': 1.2}] kwargs_fixed_source = [{ 'R_sersic': 0.6, 'n_sersic': 3, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1 }] kwargs_fixed_lens_light = [{ 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0 }] self.param_class = Param( kwargs_model, kwargs_fixed_lens=kwargs_fixed_lens, kwargs_fixed_source=kwargs_fixed_source, kwargs_fixed_lens_light=kwargs_fixed_lens_light, kwargs_lower_lens=kwargs_lower_lens, kwargs_upper_lens=kwargs_upper_lens, **kwargs_constraints) self.Likelihood = LikelihoodModule(kwargs_data_joint=kwargs_data_joint, kwargs_model=kwargs_model, param_class=self.param_class, **kwargs_likelihood) prior_means = self.param_class.kwargs2args( kwargs_lens=self.kwargs_lens, kwargs_source=self.kwargs_source, kwargs_lens_light=self.kwargs_lens_light) prior_sigmas = np.ones_like(prior_means) * 0.1 self.sampler = NestedSampler(self.Likelihood, 'gaussian', prior_means, prior_sigmas, 0.5, 0.5)
def setup(self): # data specifics sigma_bkg = 0.05 # background noise per pixel exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 10 # cutout pixel size deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2) fwhm = 0.5 # full width half max of PSF # PSF specification self.kwargs_data = sim_util.data_configure_simple( numPix, deltaPix, exp_time, sigma_bkg) data_class = ImageData(**self.kwargs_data) kwargs_psf_gaussian = { 'psf_type': 'GAUSSIAN', 'fwhm': fwhm, 'pixel_size': deltaPix, 'truncation': 3 } psf_gaussian = PSF(**kwargs_psf_gaussian) self.kwargs_psf = { 'psf_type': 'PIXEL', 'kernel_point_source': psf_gaussian.kernel_point_source, 'psf_error_map': np.zeros_like(psf_gaussian.kernel_point_source) } psf_class = PSF(**self.kwargs_psf) # 'EXTERNAL_SHEAR': external shear kwargs_shear = { 'gamma1': 0.01, 'gamma2': 0.01 } # gamma_ext: shear strength, psi_ext: shear angel (in radian) kwargs_spemd = { 'theta_E': 1., 'gamma': 1.8, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1 } lens_model_list = ['SPEP', 'SHEAR'] self.kwargs_lens = [kwargs_spemd, kwargs_shear] lens_model_class = LensModel(lens_model_list=lens_model_list) kwargs_sersic = { 'amp': 1., 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0 } # 'SERSIC_ELLIPSE': elliptical Sersic profile kwargs_sersic_ellipse = { 'amp': 1., 'R_sersic': .6, 'n_sersic': 3, 'center_x': 0, 'center_y': 0, 'e1': 0.1, 'e2': 0.1 } lens_light_model_list = ['SERSIC'] self.kwargs_lens_light = [kwargs_sersic] lens_light_model_class = LightModel( light_model_list=lens_light_model_list) source_model_list = ['SERSIC_ELLIPSE'] self.kwargs_source = [kwargs_sersic_ellipse] source_model_class = LightModel(light_model_list=source_model_list) self.kwargs_ps = [ { 'ra_source': 0.0, 'dec_source': 0.0, 'source_amp': 1. } ] # quasar point source position in the source plane and intrinsic brightness point_source_list = ['SOURCE_POSITION'] point_source_class = PointSource( point_source_type_list=point_source_list, fixed_magnification_list=[True]) kwargs_numerics = { 'supersampling_factor': 1, 'supersampling_convolution': False, 'compute_mode': 'regular', 'point_source_supersampling_factor': 1 } imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps) data_class.update_data(image_sim) self.data_class = data_class self.psf_class = psf_class self.kwargs_data['image_data'] = image_sim self.kwargs_model = { 'lens_model_list': lens_model_list, 'source_light_model_list': source_model_list, 'lens_light_model_list': lens_light_model_list, 'point_source_model_list': point_source_list, 'fixed_magnification_list': [False], 'index_lens_model_list': [[0, 1]], } self.kwargs_numerics = kwargs_numerics num_source_model = len(source_model_list) self.kwargs_constraints = { 'num_point_source_list': [4], 'image_plane_source_list': [False] * num_source_model, 'solver_type': 'NONE', # 'PROFILE', 'PROFILE_SHEAR', 'ELLIPSE', 'CENTER' } self.kwargs_likelihood = { 'force_no_add_image': True, 'source_marg': True, 'linear_prior': [1], 'image_position_uncertainty': 0.004, 'check_matched_source_position': False, 'source_position_tolerance': 0.001, 'source_position_sigma': 0.001, 'check_positive_flux': True, }
def lens_model_plot(ax, lensModel, kwargs_lens, numPix=500, deltaPix=0.01, sourcePos_x=0, sourcePos_y=0, point_source=False, with_caustics=False, with_convergence=True, coord_center_ra=0, coord_center_dec=0, coord_inverse=False, fast_caustic=True, **kwargs): """ plots a lens model (convergence) and the critical curves and caustics :param ax: matplotlib axis instance :param lensModel: LensModel() class instance :param kwargs_lens: lens model keyword argument list :param numPix: total nnumber of pixels (for convergence map) :param deltaPix: width of pixel (total frame size is deltaPix x numPix) :param sourcePos_x: float, x-position of point source (image positions computed by the lens equation) :param sourcePos_y: float, y-position of point source (image positions computed by the lens equation) :param point_source: bool, if True, illustrates and computes the image positions of the point source :param with_caustics: bool, if True, illustrates the critical curve and caustics of the system :param with_convergence: bool, if True, illustrates the convergence map :param coord_center_ra: float, x-coordinate of the center of the frame :param coord_center_dec: float, y-coordinate of the center of the frame :param coord_inverse: bool, if True, inverts the x-coordinates to go from right-to-left (effectively the RA definition) :param fast_caustic: boolean, if True, uses faster but less precise caustic calculation (might have troubles for the outer caustic (inner critical curve) :param with_convergence: boolean, if True, plots the convergence of the deflector :return: """ kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, center_ra=coord_center_ra, center_dec=coord_center_dec, inverse=coord_inverse) data = ImageData(**kwargs_data) _coords = data _frame_size = numPix * deltaPix ra0, dec0 = data.radec_at_xy_0 if coord_inverse: extent = [ra0, ra0 - _frame_size, dec0, dec0 + _frame_size] else: extent = [ra0, ra0 + _frame_size, dec0, dec0 + _frame_size] if with_convergence: kwargs_convergence = kwargs.get('kwargs_convergence', {}) convergence_plot(ax, pixel_grid=_coords, lens_model=lensModel, kwargs_lens=kwargs_lens, extent=extent, **kwargs_convergence) if with_caustics is True: kwargs_caustics = kwargs.get('kwargs_caustics', {}) caustics_plot(ax, pixel_grid=_coords, lens_model=lensModel, kwargs_lens=kwargs_lens, fast_caustic=fast_caustic, coord_inverse=coord_inverse, pixel_offset=True, **kwargs_caustics) if point_source: kwargs_point_source = kwargs.get('kwargs_point_source', {}) point_source_plot(ax, pixel_grid=_coords, lens_model=lensModel, kwargs_lens=kwargs_lens, source_x=sourcePos_x, source_y=sourcePos_y, **kwargs_point_source) if coord_inverse: ax.set_xlim([ra0, ra0 - _frame_size]) else: ax.set_xlim([ra0, ra0 + _frame_size]) ax.set_ylim([dec0, dec0 + _frame_size]) # ax.get_xaxis().set_visible(False) # ax.get_yaxis().set_visible(False) ax.autoscale(False) return ax
def arrival_time_surface(ax, lensModel, kwargs_lens, numPix=500, deltaPix=0.01, sourcePos_x=0, sourcePos_y=0, with_caustics=False, point_source=False, n_levels=10, kwargs_contours=None, image_color_list=None, letter_font_size=20): """ :param ax: matplotlib axis instance :param lensModel: LensModel() class instance :param kwargs_lens: lens model keyword argument list :param numPix: :param deltaPix: :param sourcePos_x: :param sourcePos_y: :param with_caustics: :return: """ kwargs_data = sim_util.data_configure_simple(numPix, deltaPix) data = ImageData(**kwargs_data) ra0, dec0 = data.radec_at_xy_0 origin = [ra0, dec0] _frame_size = numPix * deltaPix _coords = data x_grid, y_grid = data.pixel_coordinates lensModelExt = LensModelExtensions(lensModel) #ra_crit_list, dec_crit_list, ra_caustic_list, dec_caustic_list = lensModelExt.critical_curve_caustics( # kwargs_lens, compute_window=_frame_size, grid_scale=deltaPix/2.) x_grid1d = util.image2array(x_grid) y_grid1d = util.image2array(y_grid) fermat_surface = lensModel.fermat_potential(x_grid1d, y_grid1d, kwargs_lens, sourcePos_x, sourcePos_y) fermat_surface = util.array2image(fermat_surface) if kwargs_contours is None: kwargs_contours = {} #, cmap='Greys', vmin=-1, vmax=1) #, cmap=self._cmap, vmin=v_min, vmax=v_max) if with_caustics is True: ra_crit_list, dec_crit_list = lensModelExt.critical_curve_tiling( kwargs_lens, compute_window=_frame_size, start_scale=deltaPix / 5, max_order=10) ra_caustic_list, dec_caustic_list = lensModel.ray_shooting( ra_crit_list, dec_crit_list, kwargs_lens) plot_util.plot_line_set(ax, _coords, ra_caustic_list, dec_caustic_list, origin=origin, color='g') plot_util.plot_line_set(ax, _coords, ra_crit_list, dec_crit_list, origin=origin, color='r') if point_source is True: from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver solver = LensEquationSolver(lensModel) theta_x, theta_y = solver.image_position_from_source( sourcePos_x, sourcePos_y, kwargs_lens, min_distance=deltaPix, search_window=deltaPix * numPix) fermat_pot_images = lensModel.fermat_potential(theta_x, theta_y, kwargs_lens) _ = ax.contour( x_grid, y_grid, fermat_surface, origin='lower', # extent=[0, _frame_size, 0, _frame_size], levels=np.sort(fermat_pot_images), **kwargs_contours) # mag_images = lensModel.magnification(theta_x, theta_y, kwargs_lens) x_image, y_image = _coords.map_coord2pix(theta_x, theta_y) abc_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K'] for i in range(len(x_image)): x_ = (x_image[i] + 0.5) * deltaPix - _frame_size / 2 y_ = (y_image[i] + 0.5) * deltaPix - _frame_size / 2 if image_color_list is None: color = 'k' else: color = image_color_list[i] ax.plot(x_, y_, 'x', markersize=10, alpha=1, color=color ) # markersize=8*(1 + np.log(np.abs(mag_images[i]))) ax.text(x_ + deltaPix, y_ + deltaPix, abc_list[i], fontsize=letter_font_size, color='k') x_source, y_source = _coords.map_coord2pix(sourcePos_x, sourcePos_y) ax.plot((x_source + 0.5) * deltaPix - _frame_size / 2, (y_source + 0.5) * deltaPix - _frame_size / 2, '*k', markersize=20) else: vmin = np.min(fermat_surface) vmax = np.max(fermat_surface) levels = np.linspace(start=vmin, stop=vmax, num=n_levels) im = ax.contour( x_grid, y_grid, fermat_surface, origin='lower', # extent=[0, _frame_size, 0, _frame_size], levels=levels, **kwargs_contours) ax.get_xaxis().set_visible(False) ax.get_yaxis().set_visible(False) ax.autoscale(False) return ax
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
# data specifics sigma_bkg = 10.0 # background noise per pixel (Gaussian) exp_time = 100. # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 100 # cutout pixel size deltaPix = 0.263 # pixel size in arcsec (area per pixel = deltaPix**2) fwhm = 1.0 # full width half max of PSF (only valid when psf_type='gaussian') psf_type = 'GAUSSIAN' # 'gaussian', 'pixel', 'NONE' kernel_size = 91 # initial input simulation # generate the coordinate grid and image properties kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg) data_class = ImageData(**kwargs_data) kwargs_psf = {'psf_type': 'GAUSSIAN', 'fwhm': fwhm, 'pixel_size': deltaPix, 'truncation': 3} psf_class = PSF(**kwargs_psf) # In[8]: if __name__ == "__main__": print("simulation started") show_img = True IsTrain = True #False num_samples = 10 root_folder = "./test_sims_full_band/" if not os.path.exists(root_folder): os.mkdir(root_folder)
def setup(self): # data specifics sigma_bkg = 0.01 # background noise per pixel exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 100 # cutout pixel size deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2) fwhm = 0.3 # full width half max of PSF # PSF specification kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg) data_class = ImageData(**kwargs_data) sigma = util.fwhm2sigma(fwhm) x_grid, y_grid = util.make_grid(numPix=31, deltapix=0.05) from lenstronomy.LightModel.Profiles.gaussian import Gaussian gaussian = Gaussian() kernel_point_source = gaussian.function(x_grid, y_grid, amp=1., sigma=sigma, center_x=0, center_y=0) kernel_point_source /= np.sum(kernel_point_source) kernel_point_source = util.array2image(kernel_point_source) psf_error_map = np.zeros_like(kernel_point_source) self.kwargs_psf = { 'psf_type': 'PIXEL', 'kernel_point_source': kernel_point_source, 'psf_error_map': psf_error_map } psf_class = PSF(**self.kwargs_psf) # 'EXERNAL_SHEAR': external shear kwargs_shear = { 'gamma1': 0.01, 'gamma2': 0.01 } # gamma_ext: shear strength, psi_ext: shear angel (in radian) phi, q = 0.2, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi, q) kwargs_spemd = { 'theta_E': 1., 'gamma': 1.8, 'center_x': 0, 'center_y': 0, 'e1': e1, 'e2': e2 } lens_model_list = ['SPEP', 'SHEAR'] self.kwargs_lens = [kwargs_spemd, kwargs_shear] lens_model_class = LensModel(lens_model_list=lens_model_list) # list of light profiles (for lens and source) # 'SERSIC': spherical Sersic profile kwargs_sersic = { 'amp': 1., 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0 } # 'SERSIC_ELLIPSE': elliptical Sersic profile phi, q = 0.2, 0.9 e1, e2 = param_util.phi_q2_ellipticity(phi, q) kwargs_sersic_ellipse = { 'amp': 1., 'R_sersic': .6, 'n_sersic': 7, 'center_x': 0, 'center_y': 0, 'e1': e1, 'e2': e2 } lens_light_model_list = ['SERSIC'] self.kwargs_lens_light = [kwargs_sersic] lens_light_model_class = LightModel( light_model_list=lens_light_model_list) source_model_list = ['SERSIC_ELLIPSE'] self.kwargs_source = [kwargs_sersic_ellipse] source_model_class = LightModel(light_model_list=source_model_list) self.kwargs_ps = [ { 'ra_source': 0.0, 'dec_source': 0.0, 'source_amp': 10. } ] # quasar point source position in the source plane and intrinsic brightness point_source_class = PointSource( point_source_type_list=['SOURCE_POSITION'], fixed_magnification_list=[True]) kwargs_numerics = { 'supersampling_factor': 3, 'supersampling_convolution': False, 'compute_mode': 'regular', 'point_source_supersampling_factor': 3 } imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps) data_class.update_data(image_sim) self.imageModel = ImageLinearFit(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) self.psf_fitting = PsfFitting(self.imageModel) self.kwargs_params = { 'kwargs_lens': self.kwargs_lens, 'kwargs_source': self.kwargs_source, 'kwargs_lens_light': self.kwargs_lens_light, 'kwargs_ps': self.kwargs_ps }
def setup(self): # data specifics sigma_bkg = .05 # background noise per pixel exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) self.num_pix = 100 # cutout pixel size delta_pix = 0.05 # pixel size in arcsec (area per pixel = delta_pix**2) fwhm = 0.5 # full width half max of PSF # supersampling factor for source plane self.subgrid_res_source = 1 self.num_pix_source = self.num_pix * self.subgrid_res_source # wavelets scales for lens and source self.n_scales_source = 4 self.n_scales_lens = 3 # prepare data simulation kwargs_data = sim_util.data_configure_simple(self.num_pix, delta_pix, exp_time, sigma_bkg, inverse=True) data_class = ImageData(**kwargs_data) # generate sa pixelated gaussian PSF kernel kwargs_psf = { 'psf_type': 'GAUSSIAN', 'fwhm': fwhm, 'truncation': 5, 'pixel_size': delta_pix } psf_class = PSF(**kwargs_psf) kernel = psf_class.kernel_point_source kwargs_psf = { 'psf_type': 'PIXEL', 'kernel_point_source': kernel, 'psf_error_map': np.ones_like(kernel) * 0.001 } psf_class = PSF(**kwargs_psf) # 'EXERNAL_SHEAR': external shear kwargs_shear = { 'gamma1': 0.01, 'gamma2': 0.01 } # gamma_ext: shear strength, psi_ext: shear angel (in radian) phi, q = 0.2, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi, q) kwargs_spemd = { 'theta_E': 1., 'gamma': 1.8, 'center_x': 0, 'center_y': 0, 'e1': e1, 'e2': e2 } lens_model_list = ['SPEP', 'SHEAR'] self.kwargs_lens = [kwargs_spemd, kwargs_shear] self.lens_model_class = LensModel(lens_model_list=lens_model_list) # list of light profiles (for lens and source) # 'SERSIC': spherical Sersic profile kwargs_sersic = { 'amp': 1., 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0 } # 'SERSIC_ELLIPSE': elliptical Sersic profile phi, q = 0.2, 0.9 e1, e2 = param_util.phi_q2_ellipticity(phi, q) kwargs_sersic_ellipse = { 'amp': 1., 'R_sersic': .6, 'n_sersic': 7, 'center_x': 0, 'center_y': 0, 'e1': e1, 'e2': e2 } lens_light_model_list = ['SERSIC'] kwargs_lens_light = [kwargs_sersic] lens_light_model_class = LightModel( light_model_list=lens_light_model_list) source_model_list = ['SERSIC_ELLIPSE'] kwargs_source = [kwargs_sersic_ellipse] source_model_class = LightModel(light_model_list=source_model_list) # list of lens light profiles point_source_class = PointSource(['LENSED_POSITION']) lens_eq_solver = LensEquationSolver(lensModel=self.lens_model_class) ra_image, dec_image = lens_eq_solver.image_position_from_source( sourcePos_x=kwargs_source[0]['center_x'], sourcePos_y=kwargs_source[0]['center_y'], kwargs_lens=self.kwargs_lens) point_amp = np.ones_like(ra_image) kwargs_ps = [{ 'ra_image': ra_image, 'dec_image': dec_image, 'point_amp': point_amp }] # simulate data kwargs_numerics = {'supersampling_factor': 1} imageModel = ImageModel(data_class, psf_class, self.lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) self.image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps) data_class.update_data(self.image_sim) # retrieve the point source data only (for initial guess for source+PS solver) self.ps_sim = imageModel.image(self.kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps, source_add=False, lens_light_add=False, point_source_add=True) # define some mask self.likelihood_mask = np.zeros((self.num_pix, self.num_pix)) self.likelihood_mask[5:-5, 5:-5] = 1 # get a numerics classes numerics = NumericsSubFrame(pixel_grid=data_class, psf=psf_class) source_numerics = NumericsSubFrame( pixel_grid=data_class, psf=psf_class, supersampling_factor=self.subgrid_res_source) self.num_iter_source = 20 self.num_iter_lens = 10 self.num_iter_global = 7 self.num_iter_weights = 2 # source grid offsets self.kwargs_special = { 'delta_x_source_grid': 0, 'delta_y_source_grid': 0, } # init the solvers # SOLVER SOURCE, with analysis formulation self.source_model_class = LightModel(['SLIT_STARLETS']) self.kwargs_source = [{'n_scales': self.n_scales_source}] self.solver_source_ana = SparseSolverSource( data_class, self.lens_model_class, numerics, source_numerics, self.source_model_class, source_interpolation='bilinear', minimal_source_plane=False, use_mask_for_minimal_source_plane=True, min_num_pix_source=20, sparsity_prior_norm=1, force_positivity=True, formulation='analysis', verbose=False, show_steps=False, min_threshold=5, threshold_increment_high_freq=1, threshold_decrease_type='exponential', num_iter_source=self.num_iter_source, num_iter_weights=self.num_iter_weights) self.solver_source_ana.set_likelihood_mask(self.likelihood_mask) # SOLVER SOURCE + LENS, with synthesis formulation self.lens_light_model_class = LightModel(['SLIT_STARLETS']) self.kwargs_lens_light = [{'n_scales': self.n_scales_lens}] self.solver_lens_syn = SparseSolverSourceLens( data_class, self.lens_model_class, numerics, source_numerics, self.source_model_class, self.lens_light_model_class, source_interpolation='bilinear', minimal_source_plane=False, use_mask_for_minimal_source_plane=True, min_num_pix_source=20, sparsity_prior_norm=1, force_positivity=True, formulation='synthesis', verbose=False, show_steps=False, min_threshold=3, threshold_increment_high_freq=1, threshold_decrease_type='linear', num_iter_global=self.num_iter_global, num_iter_source=self.num_iter_source, num_iter_lens=self.num_iter_lens, num_iter_weights=self.num_iter_weights) self.solver_lens_syn.set_likelihood_mask(self.likelihood_mask)
def lens_model_plot_custom(image, ax, lensModel, kwargs_lens, numPix=500, deltaPix=0.01, sourcePos_x=0, sourcePos_y=0, point_source=False, with_caustics=False): """ Overlay the critical curves and caustics over the provided lensed image Parameters ---------- image : np.array ax : matplotlib.pyplot.axes """ kwargs_data = sim_util.data_configure_simple(numPix, deltaPix) data = ImageData(**kwargs_data) _coords = data _frame_size = numPix * deltaPix x_grid, y_grid = data.pixel_coordinates lensModelExt = LensModelExtensions(lensModel) _ = ax.matshow(image, origin='lower', extent=[0, _frame_size, 0, _frame_size]) if with_caustics is True: ra_crit_list, dec_crit_list = lensModelExt.critical_curve_tiling( kwargs_lens, compute_window=_frame_size, start_scale=deltaPix, max_order=20) ra_caustic_list, dec_caustic_list = lensModel.ray_shooting( ra_crit_list, dec_crit_list, kwargs_lens) plot_util.plot_line_set(ax, _coords, ra_caustic_list, dec_caustic_list, color='y') plot_util.plot_line_set(ax, _coords, ra_crit_list, dec_crit_list, color='r') if point_source: solver = LensEquationSolver(lensModel) theta_x, theta_y = solver.image_position_from_source( sourcePos_x, sourcePos_y, kwargs_lens, min_distance=deltaPix, search_window=deltaPix * numPix) mag_images = lensModel.magnification(theta_x, theta_y, kwargs_lens) x_image, y_image = _coords.map_coord2pix(theta_x, theta_y) abc_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K'] for i in range(len(x_image)): x_ = (x_image[i] + 0.5) * deltaPix y_ = (y_image[i] + 0.5) * deltaPix ax.plot(x_, y_, 'dk', markersize=4 * (1 + np.log(np.abs(mag_images[i]))), alpha=0.5) ax.text(x_, y_, abc_list[i], fontsize=20, color='k') x_source, y_source = _coords.map_coord2pix(sourcePos_x, sourcePos_y) ax.plot((x_source + 0.5) * deltaPix, (y_source + 0.5) * deltaPix, '*r', markersize=5) #ax.plot(numPix * deltaPix*0.5 + pred['lens_mass_center_x'] + pred['src_light_center_x'], numPix * deltaPix*0.5 + pred['lens_mass_center_y'] + pred['src_light_center_y'], '*k', markersize=5) ax.get_xaxis().set_visible(False) ax.get_yaxis().set_visible(False) ax.autoscale(False) return ax
delta_t.to_csv("../data/time_delays.csv") # data specifics for the lens image sigma_bkg = .05 # background noise per pixel (Gaussian) exp_time = 1. # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) delta_pix_x = np.abs(lens_sky.ra[0] - lens_sky.ra[1]).to(u.arcsec).value delta_pix_y = np.abs(lens_sky.dec[0] - lens_sky.dec[M]).to(u.arcsec).value kwargs_data = { "image_data": im, "exposure_time": exp_time, "background_rms": sigma_bkg, "transform_pix2angle": np.array([[-delta_pix_x, 0], [0, delta_pix_y]]), "ra_at_xy_0": 0, "dec_at_xy_0": 0 } data_class = ImageData(**kwargs_data) x = lens_sky.ra.to(u.arcsec) - lens_sky.ra.max() y = lens_sky.dec.to(u.arcsec) - lens_sky.dec.min() x_center = (x[M // 2]).to(u.arcsec).value y_center = (y[M * N // 2]).to(u.arcsec).value grid = np.column_stack([x, y]) ee = 0.611 # encircled energy, source https://www.stsci.edu/hst/instrumentation/acs/data-analysis/aperture-corrections # collect 4 pixels closest to the image position (correspond to 0.1 arcsec since pixelscale is 0.05 arcsec) f = [] A = [] for i in range(4): pos = image_position.to_numpy()[i] diff = np.square(pos - grid.value) dist = np.einsum("ij -> i", diff) indexes = np.argpartition(dist, kth=4)[:4]
def distortions(lensModel, kwargs_lens, num_pix=100, delta_pix=0.05, center_ra=0, center_dec=0, differential_scale=0.0001, smoothing_scale=None, **kwargs): """ :param lensModel: LensModel instance :param kwargs_lens: lens model keyword argument list :param num_pix: number of pixels per axis :param delta_pix: pixel scale per axis :param center_ra: center of the grid :param center_dec: center of the grid :param differential_scale: scale of the finite derivative length in units of angles :param smoothing_scale: float or None, Gaussian FWHM of a smoothing kernel applied before plotting :return: matplotlib instance with different panels """ kwargs_grid = sim_util.data_configure_simple(num_pix, delta_pix, center_ra=center_ra, center_dec=center_dec) _coords = ImageData(**kwargs_grid) _frame_size = num_pix * delta_pix ra_grid, dec_grid = _coords.pixel_coordinates extensions = LensModelExtensions(lensModel=lensModel) ra_grid1d = util.image2array(ra_grid) dec_grid1d = util.image2array(dec_grid) lambda_rad, lambda_tan, orientation_angle, dlambda_tan_dtan, dlambda_tan_drad, dlambda_rad_drad, dlambda_rad_dtan, dphi_tan_dtan, dphi_tan_drad, dphi_rad_drad, dphi_rad_dtan = extensions.radial_tangential_differentials( ra_grid1d, dec_grid1d, kwargs_lens=kwargs_lens, center_x=center_ra, center_y=center_dec, smoothing_3rd=differential_scale, smoothing_2nd=None) lambda_rad2d, lambda_tan2d, orientation_angle2d, dlambda_tan_dtan2d, dlambda_tan_drad2d, dlambda_rad_drad2d, dlambda_rad_dtan2d, dphi_tan_dtan2d, dphi_tan_drad2d, dphi_rad_drad2d, dphi_rad_dtan2d = util.array2image(lambda_rad), \ util.array2image(lambda_tan), util.array2image(orientation_angle), util.array2image(dlambda_tan_dtan), util.array2image(dlambda_tan_drad), util.array2image(dlambda_rad_drad), util.array2image(dlambda_rad_dtan), \ util.array2image(dphi_tan_dtan), util.array2image(dphi_tan_drad), util.array2image(dphi_rad_drad), util.array2image(dphi_rad_dtan) if smoothing_scale is not None: lambda_rad2d = ndimage.gaussian_filter(lambda_rad2d, sigma=smoothing_scale / delta_pix) dlambda_rad_drad2d = ndimage.gaussian_filter(dlambda_rad_drad2d, sigma=smoothing_scale / delta_pix) lambda_tan2d = np.abs(lambda_tan2d) # the magnification cut is made to make a stable integral/convolution lambda_tan2d[lambda_tan2d > 100] = 100 lambda_tan2d = ndimage.gaussian_filter(lambda_tan2d, sigma=smoothing_scale / delta_pix) # the magnification cut is made to make a stable integral/convolution dlambda_tan_dtan2d[dlambda_tan_dtan2d > 100] = 100 dlambda_tan_dtan2d[dlambda_tan_dtan2d < -100] = -100 dlambda_tan_dtan2d = ndimage.gaussian_filter(dlambda_tan_dtan2d, sigma=smoothing_scale / delta_pix) orientation_angle2d = ndimage.gaussian_filter(orientation_angle2d, sigma=smoothing_scale / delta_pix) dphi_tan_dtan2d = ndimage.gaussian_filter(dphi_tan_dtan2d, sigma=smoothing_scale / delta_pix) def _plot_frame(ax, map, vmin, vmax, text_string): """ :param ax: matplotlib.axis instance :param map: 2d array :param vmin: minimum plotting scale :param vmax: maximum plotting scale :param text_string: string to describe the label :return: """ font_size = 10 _arrow_size = 0.02 im = ax.matshow(map, extent=[0, _frame_size, 0, _frame_size], vmin=vmin, vmax=vmax) ax.get_xaxis().set_visible(False) ax.get_yaxis().set_visible(False) ax.autoscale(False) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) cb = plt.colorbar(im, cax=cax, orientation='vertical') #cb.set_label(text_string, fontsize=10) #plot_util.scale_bar(ax, _frame_size, dist=1, text='1"', font_size=font_size) plot_util.text_description(ax, _frame_size, text=text_string, color="k", backgroundcolor='w', font_size=font_size) #if 'no_arrow' not in kwargs or not kwargs['no_arrow']: # plot_util.coordinate_arrows(ax, _frame_size, _coords, # color='w', arrow_size=_arrow_size, # font_size=font_size) f, axes = plt.subplots(3, 4, figsize=(12, 8)) _plot_frame(axes[0, 0], lambda_rad2d, vmin=0.6, vmax=1.4, text_string=r"$\lambda_{rad}$") _plot_frame(axes[0, 1], lambda_tan2d, vmin=-20, vmax=20, text_string=r"$\lambda_{tan}$") _plot_frame(axes[0, 2], orientation_angle2d, vmin=-np.pi / 10, vmax=np.pi / 10, text_string=r"$\phi$") _plot_frame(axes[0, 3], util.array2image(lambda_tan * lambda_rad), vmin=-20, vmax=20, text_string='magnification') _plot_frame(axes[1, 0], dlambda_rad_drad2d / lambda_rad2d, vmin=-.1, vmax=.1, text_string='dlambda_rad_drad') _plot_frame(axes[1, 1], dlambda_tan_dtan2d / lambda_tan2d, vmin=-20, vmax=20, text_string='dlambda_tan_dtan') _plot_frame(axes[1, 2], dlambda_tan_drad2d / lambda_tan2d, vmin=-20, vmax=20, text_string='dlambda_tan_drad') _plot_frame(axes[1, 3], dlambda_rad_dtan2d / lambda_rad2d, vmin=-.1, vmax=.1, text_string='dlambda_rad_dtan') _plot_frame(axes[2, 0], dphi_rad_drad2d, vmin=-.1, vmax=.1, text_string='dphi_rad_drad') _plot_frame(axes[2, 1], dphi_tan_dtan2d, vmin=0, vmax=20, text_string='dphi_tan_dtan: curvature radius') _plot_frame(axes[2, 2], dphi_tan_drad2d, vmin=-.1, vmax=.1, text_string='dphi_tan_drad') _plot_frame(axes[2, 3], dphi_rad_dtan2d, vmin=0, vmax=20, text_string='dphi_rad_dtan') return f, axes
def setup(self): # data specifics sigma_bkg = .05 # background noise per pixel exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 100 # cutout pixel size deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2) fwhm = 0.5 # full width half max of PSF # PSF specification kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg, inverse=True) data_class = ImageData(**kwargs_data) kwargs_psf = { 'psf_type': 'GAUSSIAN', 'fwhm': fwhm, 'truncation': 5, 'pixel_size': deltaPix } psf_class = PSF(**kwargs_psf) kernel = psf_class.kernel_point_source kwargs_psf = { 'psf_type': 'PIXEL', 'kernel_point_source': kernel, 'psf_error_map': np.ones_like(kernel) * 0.001 } psf_class = PSF(**kwargs_psf) # 'EXERNAL_SHEAR': external shear kwargs_shear = { 'gamma1': 0.01, 'gamma2': 0.01 } # gamma_ext: shear strength, psi_ext: shear angel (in radian) phi, q = 0.2, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi, q) kwargs_spemd = { 'theta_E': 1., 'gamma': 1.8, 'center_x': 0, 'center_y': 0, 'e1': e1, 'e2': e2 } lens_model_list = ['SPEP', 'SHEAR'] self.kwargs_lens = [kwargs_spemd, kwargs_shear] lens_model_class = LensModel(lens_model_list=lens_model_list) # list of light profiles (for lens and source) # 'SERSIC': spherical Sersic profile kwargs_sersic = { 'amp': 1., 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0 } # 'SERSIC_ELLIPSE': elliptical Sersic profile phi, q = 0.2, 0.9 e1, e2 = param_util.phi_q2_ellipticity(phi, q) kwargs_sersic_ellipse = { 'amp': 1., 'R_sersic': .6, 'n_sersic': 7, 'center_x': 0, 'center_y': 0, 'e1': e1, 'e2': e2 } lens_light_model_list = ['SERSIC'] self.kwargs_lens_light = [kwargs_sersic] lens_light_model_class = LightModel( light_model_list=lens_light_model_list) source_model_list = ['SERSIC_ELLIPSE'] self.kwargs_source = [kwargs_sersic_ellipse] source_model_class = LightModel(light_model_list=source_model_list) self.kwargs_ps = [ { 'ra_source': 0.01, 'dec_source': 0.0, 'source_amp': 1. } ] # quasar point source position in the source plane and intrinsic brightness point_source_class = PointSource( point_source_type_list=['SOURCE_POSITION'], fixed_magnification_list=[True]) kwargs_numerics = { 'supersampling_factor': 2, 'supersampling_convolution': False } imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps) data_class.update_data(image_sim) self.imageModel = ImageLinearFit(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) self.solver = LensEquationSolver(lensModel=self.imageModel.LensModel)
def test_update_data(self): kwargs_data = sim_util.data_configure_simple(numPix=10, deltaPix=1, exposure_time=1, background_rms=1, inverse=True) data_class = ImageData(**kwargs_data) self.imageModel.update_data(data_class) assert self.imageModel.Data.num_pixel == 100