def test_multinest(self): # Nested sampler tests # further decrease the parameter space for nested samplers to run faster fittingSequence = FittingSequence(self.kwargs_data_joint, self.kwargs_model, self.kwargs_constraints, self.kwargs_likelihood, self.kwargs_params) fitting_list = [] kwargs_update = { 'ps_add_fixed': [[0, ['ra_source', 'dec_source'], [0, 0]]], 'lens_light_add_fixed': [[ 0, ['n_sersic', 'R_sersic', 'center_x', 'center_y'], [4, .1, 0, 0] ]], 'source_add_fixed': [[ 0, ['R_sersic', 'e1', 'e2', 'center_x', 'center_y'], [.6, .1, .1, 0, 0] ]], 'lens_add_fixed': [[ 0, ['gamma', 'theta_E', 'e1', 'e2', 'center_x', 'center_y'], [1.8, 1., .1, .1, 0, 0] ], [1, ['gamma1', 'gamma2'], [0.01, 0.01]]], 'change_source_lower_limit': [[0, ['n_sersic'], [2.9]]], 'change_source_upper_limit': [[0, ['n_sersic'], [3.1]]] } fitting_list.append(['update_settings', kwargs_update]) kwargs_multinest = { 'sampler_type': 'MULTINEST', 'kwargs_run': { 'n_live_points': 10, 'evidence_tolerance': 0.5, 'sampling_efficiency': 0.8, # 1 for posterior-only, 0 for evidence-only 'importance_nested_sampling': False, 'multimodal': True, 'const_efficiency_mode': False, # reduce sampling_efficiency to 5% when True }, 'remove_output_dir': True, } fitting_list.append(['nested_sampling', kwargs_multinest]) chain_list2 = fittingSequence.fit_sequence(fitting_list) kwargs_fixed = fittingSequence._updateManager.fixed_kwargs npt.assert_almost_equal(kwargs_fixed[0][1]['gamma1'], 0.01, decimal=2) assert fittingSequence._updateManager._lower_kwargs[1][0][ 'n_sersic'] == 2.9 assert fittingSequence._updateManager._upper_kwargs[1][0][ 'n_sersic'] == 3.1 kwargs_test = {'kwargs_lens': 1} fittingSequence.update_state(kwargs_test) kwargs_out = fittingSequence.best_fit(bijective=True) assert kwargs_out['kwargs_lens'] == 1
def run_mcmc(self, mcmc_numerics): """Sample from the joint likelihood """ fitting_seq = FittingSequence(self.kwargs_data_joint, self.kwargs_model, self.kwargs_constraints, self.kwargs_likelihood, self.kwargs_params) fitting_kwargs_list = [['MCMC', mcmc_numerics]] #with script_utils.HiddenPrints(): chain_list_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result_mcmc = fitting_seq.best_fit() return chain_list_mcmc, kwargs_result_mcmc
def make_lensmodel(lens_info, theta_E, source_info, box_f): # lens data specifics lens_image = lens_info['image'] psf_lens = lens_info['psf'] background_rms = background_rms_image(5, lens_image) exposure_time = 100 kwargs_data_lens = sim_util.data_configure_simple(len(lens_image), lens_info['deltapix'], exposure_time, background_rms) kwargs_data_lens['image_data'] = lens_image data_class_lens = ImageData(**kwargs_data_lens) #PSF kwargs_psf_lens = { 'psf_type': 'PIXEL', 'pixel_size': lens_info['deltapix'], 'kernel_point_source': psf_lens } psf_class_lens = PSF(**kwargs_psf_lens) # lens light model lens_light_model_list = ['SERSIC_ELLIPSE'] lens_light_model_class = LightModel(light_model_list=lens_light_model_list) kwargs_model = {'lens_light_model_list': lens_light_model_list} kwargs_numerics_galfit = {'supersampling_factor': 1} kwargs_constraints = {} kwargs_likelihood = {'check_bounds': True} image_band = [kwargs_data_lens, kwargs_psf_lens, kwargs_numerics_galfit] multi_band_list = [image_band] kwargs_data_joint = { 'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear' } # Sersic component fixed_lens_light = [{}] kwargs_lens_light_init = [{ 'R_sersic': .1, 'n_sersic': 4, 'e1': 0, 'e2': 0, 'center_x': 0, 'center_y': 0 }] kwargs_lens_light_sigma = [{ 'n_sersic': 0.5, 'R_sersic': 0.2, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1 }] kwargs_lower_lens_light = [{ 'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.01, 'n_sersic': 0.5, 'center_x': -10, 'center_y': -10 }] kwargs_upper_lens_light = [{ 'e1': 0.5, 'e2': 0.5, 'R_sersic': 10, 'n_sersic': 8, 'center_x': 10, 'center_y': 10 }] lens_light_params = [ kwargs_lens_light_init, kwargs_lens_light_sigma, fixed_lens_light, kwargs_lower_lens_light, kwargs_upper_lens_light ] kwargs_params = {'lens_light_model': lens_light_params} fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) fitting_kwargs_list = [[ 'PSO', { 'sigma_scale': 1., 'n_particles': 50, 'n_iterations': 50 } ]] chain_list = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_seq.best_fit() modelPlot = ModelPlot(multi_band_list, kwargs_model, kwargs_result) # Lens light best result kwargs_light_lens = kwargs_result['kwargs_lens_light'][0] #Lens model kwargs_lens_list = [{ 'theta_E': theta_E, 'e1': kwargs_light_lens['e1'], 'e2': kwargs_light_lens['e2'], 'center_x': kwargs_light_lens['center_x'], 'center_y': kwargs_light_lens['center_y'] }] lensModel = LensModel(['SIE']) lme = LensModelExtensions(lensModel) #random position for the source x_crit_list, y_crit_list = lme.critical_curve_tiling( kwargs_lens_list, compute_window=(len(source_info['image'])) * (source_info['deltapix']), start_scale=source_info['deltapix'], max_order=10) if len(x_crit_list) > 2 and len(y_crit_list) > 2: x_caustic_list, y_caustic_list = lensModel.ray_shooting( x_crit_list, y_crit_list, kwargs_lens_list) xsamp0 = np.arange( min(x_caustic_list) - min(x_caustic_list) * box_f[0], max(x_caustic_list) + max(x_caustic_list) * box_f[1], 0.1) xsamp = xsamp0[abs(xsamp0.round(1)) != 0.1] ysamp0 = np.arange( min(y_caustic_list) - min(y_caustic_list) * box_f[0], max(y_caustic_list) + max(y_caustic_list) * box_f[1], 0.1) ysamp = ysamp0[abs(ysamp0.round(1)) != 0.1] if len(xsamp) == 0 or len(ysamp) == 0: x_shift, y_shift = 0.15, 0.15 #arcseconds else: y_shift = rand.sample(list(ysamp), 1)[0] x_shift = rand.sample(list(xsamp), 1)[0] else: x_shift, y_shift = -0.15, 0.15 #arcseconds x_caustic_list = [0] y_caustic_list = [0] solver = LensEquationSolver(lensModel) theta_ra, theta_dec = solver.image_position_from_source( x_shift, y_shift, kwargs_lens_list) if len(theta_ra) <= 1: x_shift, y_shift = -0.2, -0.2 #arcseconds1 if abs(x_shift) >= int(theta_E) or abs(y_shift) >= int(theta_E): x_shift, y_shift = 0.3, -0.3 print('BLABLA') print('HERE', min(x_caustic_list) - min(x_caustic_list) * box_f[0], max(x_caustic_list) + max(x_caustic_list) * box_f[1], min(y_caustic_list) - min(y_caustic_list) * box_f[0], max(y_caustic_list) + max(y_caustic_list) * box_f[1]) return { 'lens_light_model_list': ['SERSIC_ELLIPSE'], 'kwargs_light_lens': [kwargs_light_lens], 'lens_light_model_class': lens_light_model_class, 'kwargs_lens_list': kwargs_lens_list, 'kwargs_data_lens': kwargs_data_lens, 'source_shift': [x_shift, y_shift] }
'n_particles': 300, 'n_iterations': 200 } ], [ 'PSO', { 'sigma_scale': 1., 'n_particles': 300, 'n_iterations': 200 } ] ] start_time = time.time() chain_list_pso = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_seq.best_fit() end_time = time.time() print(end_time - start_time, 'total time needed for computation') print( '============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ' ) #%% kwargs_result = fitting_seq.best_fit(bijective=True) from lenstronomy.Plots import chain_plot as chain_plot for i in range(len(chain_list_pso)): chain_plot.plot_chain_list(chain_list_pso, i) plt.close() #and now we run the MCMC
def main(): args = parse_args() test_cfg = TestConfig.from_file(args.test_config_file_path) train_val_cfg = TrainValConfig.from_file(test_cfg.train_val_config_file_path) baobab_cfg = get_baobab_config(test_cfg.data.test_dir) # Set device and default data type device = torch.device(test_cfg.device_type) if device.type == 'cuda': torch.set_default_tensor_type('torch.cuda.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') seed_everything(test_cfg.global_seed) ############ # Data I/O # ############ test_data = XYCosmoData(test_cfg.data.test_dir, data_cfg=train_val_cfg.data) master_truth = test_data.cosmo_df master_truth = metadata_utils.add_qphi_columns(master_truth) master_truth = metadata_utils.add_gamma_psi_ext_columns(master_truth) if test_cfg.data.lens_indices is None: if args.lens_indices_path is None: # Test on all n_test lenses in the test set n_test = test_cfg.data.n_test lens_range = range(n_test) else: # Test on the lens indices in a text file at the specified path lens_range = [] with open(args.lens_indices_path, "r") as f: for line in f: lens_range.append(int(line.strip())) n_test = len(lens_range) print("Performing H0 inference on {:d} specified lenses...".format(n_test)) else: if args.lens_indices_path is None: # Test on the lens indices specified in the test config file lens_range = test_cfg.data.lens_indices n_test = len(lens_range) print("Performing H0 inference on {:d} specified lenses...".format(n_test)) else: raise ValueError("Specific lens indices were specified in both the test config file and the command-line argument.") batch_size = max(lens_range) + 1 # Output directory into which the H0 histograms and H0 samples will be saved out_dir = test_cfg.out_dir if not os.path.exists(out_dir): os.makedirs(out_dir) print("Destination folder path: {:s}".format(out_dir)) else: raise OSError("Destination folder already exists.") ###################### # Load trained state # ###################### # Instantiate loss function, to append to the MCMC objective as the prior orig_Y_cols = train_val_cfg.data.Y_cols # Instantiate MCMC parameter penalty function params_to_remove = ['lens_light_R_sersic'] #'src_light_R_sersic'] mcmc_Y_cols = [col for col in orig_Y_cols if col not in params_to_remove] mcmc_Y_dim = len(mcmc_Y_cols) null_spread = True #init_D_dt = np.random.uniform(0.0, 10000.0, size=(batch_size, n_walkers, 1)) # FIXME: init H0 hardcoded kwargs_model = dict(lens_model_list=['PEMD', 'SHEAR'], point_source_model_list=['SOURCE_POSITION'], source_light_model_list=['SERSIC_ELLIPSE']) astro_sig = test_cfg.image_position_likelihood.sigma # Get H0 samples for each system if not test_cfg.time_delay_likelihood.baobab_time_delays: if 'abcd_ordering_i' not in master_truth: raise ValueError("If the time delay measurements were not generated using Baobab, the user must specify the order of image positions in which the time delays are listed, in order of increasing dec.") kwargs_lens_eq_solver = {'min_distance': 0.05, 'search_window': baobab_cfg.instrument.pixel_scale*baobab_cfg.image.num_pix, 'num_iter_max': 100} #n_walkers = test_cfg.numerics.mcmc.walkerRatio*(mcmc_Y_dim + 1) # BNN params + H0 times walker ratio #init_pos = np.tile(master_truth[mcmc_Y_cols].iloc[:batch_size].values[:, np.newaxis, :], [1, n_walkers, 1]) #init_D_dt = np.random.uniform(0.0, 10000.0, size=(batch_size, n_walkers, 1)) #print(init_pos.shape, init_D_dt.shape) total_progress = tqdm(total=n_test) # For each lens system... for i, lens_i in enumerate(lens_range): # Each lens gets a unique random state for td and vd measurement error realizations. rs_lens = np.random.RandomState(lens_i) ########################### # Relevant data and prior # ########################### data_i = master_truth.iloc[lens_i].copy() # Init values for the lens model params init_info = dict(zip(mcmc_Y_cols, data_i[mcmc_Y_cols].values)) # truth params lcdm = LCDM(z_lens=data_i['z_lens'], z_source=data_i['z_src'], flat=True) true_img_dec = np.array(literal_eval(data_i['y_image'])) n_img = len(true_img_dec) true_td = np.array(literal_eval(data_i['true_td'])) measured_td = true_td + rs_lens.randn(*true_td.shape)*test_cfg.error_model.time_delay_error measured_td_sig = test_cfg.time_delay_likelihood.sigma # np.ones(n_img - 1)* measured_img_dec = true_img_dec + rs_lens.randn(n_img)*astro_sig increasing_dec_i = np.argsort(true_img_dec) #np.argsort(measured_img_dec) measured_td = h0_utils.reorder_to_tdlmc(measured_td, increasing_dec_i, range(n_img)) # need to use measured dec to order measured_img_dec = h0_utils.reorder_to_tdlmc(measured_img_dec, increasing_dec_i, range(n_img)) measured_td_wrt0 = measured_td[1:] - measured_td[0] kwargs_data_joint = dict(time_delays_measured=measured_td_wrt0, time_delays_uncertainties=measured_td_sig, ) ############################# # Parameter init and bounds # ############################# lens_kwargs = mcmc_utils.get_lens_kwargs(init_info, null_spread=null_spread) ps_kwargs = mcmc_utils.get_ps_kwargs_src_plane(init_info, astro_sig, null_spread=null_spread) src_light_kwargs = mcmc_utils.get_light_kwargs(init_info['src_light_R_sersic'], null_spread=null_spread) special_kwargs = mcmc_utils.get_special_kwargs(n_img, astro_sig, D_dt_sigma=2000, null_spread=null_spread) # image position offset and time delay distance, aka the "special" parameters kwargs_params = {'lens_model': lens_kwargs, 'point_source_model': ps_kwargs, 'source_model': src_light_kwargs, 'special': special_kwargs,} if test_cfg.numerics.solver_type == 'NONE': solver_type = 'NONE' else: solver_type = 'PROFILE_SHEAR' if n_img == 4 else 'CENTER' #solver_type = 'NONE' kwargs_constraints = {'num_point_source_list': [n_img], 'Ddt_sampling': True, 'solver_type': solver_type,} kwargs_likelihood = {'time_delay_likelihood': True, 'sort_images_by_dec': True, 'prior_lens': [], 'prior_special': [], 'check_bounds': True, 'check_matched_source_position': False, 'source_position_tolerance': 0.01, 'source_position_sigma': 0.01, 'source_position_likelihood': False, 'custom_logL_addition': None, 'kwargs_lens_eq_solver': kwargs_lens_eq_solver} ########################### # MCMC posterior sampling # ########################### fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, verbose=False, mpi=False) if i == 0: param_class = fitting_seq._updateManager.param_class n_params, param_class_Y_cols = param_class.num_param() #init_pos = mcmc_utils.reorder_to_param_class(mcmc_Y_cols, param_class_Y_cols, init_pos, init_D_dt) # MCMC sample from the post-processed BNN posterior jointly with cosmology lens_i_start_time = time.time() #test_cfg.numerics.mcmc.update(init_samples=init_pos[lens_i, :, :]) fitting_kwargs_list_mcmc = [['MCMC', test_cfg.numerics.mcmc]] #with HiddenPrints(): #try: chain_list_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list_mcmc) kwargs_result_mcmc = fitting_seq.best_fit() #except: # print("lens {:d} skipped".format(lens_i)) # total_progress.update(1) # continue lens_i_end_time = time.time() inference_time = (lens_i_end_time - lens_i_start_time)/60.0 # min ############################# # Plotting the MCMC samples # ############################# # sampler_type : 'EMCEE' # samples_mcmc : np.array of shape `[n_mcmc_eval, n_params]` # param_mcmc : list of str of length n_params, the parameter names sampler_type, samples_mcmc, param_mcmc, _ = chain_list_mcmc[0] new_samples_mcmc = mcmc_utils.postprocess_mcmc_chain(kwargs_result_mcmc, samples_mcmc, kwargs_model, lens_kwargs[2], ps_kwargs[2], src_light_kwargs[2], special_kwargs[2], kwargs_constraints) # Plot D_dt histogram D_dt_samples = new_samples_mcmc['D_dt'].values true_D_dt = lcdm.D_dt(H_0=data_i['H0'], Om0=0.3) data_i['D_dt'] = true_D_dt # Export D_dt samples for this lens lens_inference_dict = dict( D_dt_samples=D_dt_samples, # kappa_ext=0 for these samples inference_time=inference_time, true_D_dt=true_D_dt, ) lens_inference_dict_save_path = os.path.join(out_dir, 'D_dt_dict_{0:04d}.npy'.format(lens_i)) np.save(lens_inference_dict_save_path, lens_inference_dict) # Optionally export the MCMC samples if test_cfg.export.mcmc_samples: mcmc_samples_path = os.path.join(out_dir, 'mcmc_samples_{0:04d}.csv'.format(lens_i)) new_samples_mcmc.to_csv(mcmc_samples_path, index=None) # Optionally export the D_dt histogram if test_cfg.export.D_dt_histogram: cleaned_D_dt_samples = h0_utils.remove_outliers_from_lognormal(D_dt_samples, 3) _ = plotting_utils.plot_D_dt_histogram(cleaned_D_dt_samples, lens_i, true_D_dt, save_dir=out_dir) # Optionally export the plot of MCMC chain if test_cfg.export.mcmc_chain: mcmc_chain_path = os.path.join(out_dir, 'mcmc_chain_{0:04d}.png'.format(lens_i)) plotting_utils.plot_mcmc_chain(chain_list_mcmc, mcmc_chain_path) # Optionally export posterior cornerplot of select lens model parameters with D_dt if test_cfg.export.mcmc_corner: mcmc_corner_path = os.path.join(out_dir, 'mcmc_corner_{0:04d}.png'.format(lens_i)) plotting_utils.plot_mcmc_corner(new_samples_mcmc[test_cfg.export.mcmc_cols], data_i[test_cfg.export.mcmc_cols], test_cfg.export.mcmc_col_labels, mcmc_corner_path) total_progress.update(1) gc.collect() total_progress.close()
def main(): args = parse_args() test_cfg = TestConfig.from_file(args.test_config_file_path) train_val_cfg = TrainValConfig.from_file( test_cfg.train_val_config_file_path) # Set device and default data type device = torch.device(test_cfg.device_type) if device.type == 'cuda': torch.set_default_tensor_type('torch.cuda.FloatTensor') else: torch.set_default_tensor_type('torch.FloatTensor') seed_everything(test_cfg.global_seed) ############ # Data I/O # ############ test_data = TDLMCData(data_cfg=train_val_cfg.data, rung_i=args.rung_idx) master_truth = test_data.cosmo_df if test_cfg.data.lens_indices is None: if args.lens_indices_path is None: # Test on all n_test lenses in the test set n_test = test_cfg.data.n_test lens_range = range(n_test) else: # Test on the lens indices in a text file at the specified path lens_range = [] with open(args.lens_indices_path, "r") as f: for line in f: lens_range.append(int(line.strip())) n_test = len(lens_range) print("Performing H0 inference on {:d} specified lenses...".format( n_test)) else: if args.lens_indices_path is None: # Test on the lens indices specified in the test config file lens_range = test_cfg.data.lens_indices n_test = len(lens_range) print("Performing H0 inference on {:d} specified lenses...".format( n_test)) else: raise ValueError( "Specific lens indices were specified in both the test config file and the command-line argument." ) batch_size = max(lens_range) + 1 test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, drop_last=True) # Output directory into which the H0 histograms and H0 samples will be saved out_dir = test_cfg.out_dir if not os.path.exists(out_dir): os.makedirs(out_dir) print("Destination folder path: {:s}".format(out_dir)) else: raise OSError("Destination folder already exists.") ###################### # Load trained state # ###################### # Instantiate loss function, to append to the MCMC objective as the prior orig_Y_cols = train_val_cfg.data.Y_cols loss_fn = getattr(h0rton.losses, train_val_cfg.model.likelihood_class)( Y_dim=train_val_cfg.data.Y_dim, device=device) # Instantiate MCMC parameter penalty function params_to_remove = ['lens_light_R_sersic'] #, 'src_light_R_sersic'] mcmc_Y_cols = [col for col in orig_Y_cols if col not in params_to_remove] mcmc_Y_dim = len(mcmc_Y_cols) mcmc_loss_fn = getattr( h0rton.losses, train_val_cfg.model.likelihood_class)( Y_dim=train_val_cfg.data.Y_dim - len(params_to_remove), device=device) remove_param_idx, remove_idx = mcmc_utils.get_idx_for_params( mcmc_loss_fn.out_dim, orig_Y_cols, params_to_remove, train_val_cfg.model.likelihood_class) mcmc_train_Y_mean = np.delete(train_val_cfg.data.train_Y_mean, remove_param_idx) mcmc_train_Y_std = np.delete(train_val_cfg.data.train_Y_std, remove_param_idx) parameter_penalty = mcmc_utils.HybridBNNPenalty( mcmc_Y_cols, train_val_cfg.model.likelihood_class, mcmc_train_Y_mean, mcmc_train_Y_std, test_cfg.h0_posterior.exclude_velocity_dispersion, device) custom_logL_addition = parameter_penalty.evaluate if test_cfg.lens_posterior_type.startswith( 'default') else None null_spread = True if test_cfg.lens_posterior_type == 'truth' else False # Instantiate model net = getattr( h0rton.models, train_val_cfg.model.architecture)(num_classes=loss_fn.out_dim) net.to(device) # Load trained weights from saved state net, epoch = train_utils.load_state_dict_test(test_cfg.state_dict_path, net, train_val_cfg.optim.n_epochs, device) with torch.no_grad(): net.eval() for X_ in test_loader: X = X_.to(device) pred = net(X) break mcmc_pred = pred.cpu().numpy() mcmc_pred = mcmc_utils.remove_parameters_from_pred(mcmc_pred, remove_idx, return_as_tensor=False) # Instantiate posterior for BNN samples, to initialize the walkers bnn_post = getattr(h0rton.h0_inference.gaussian_bnn_posterior, loss_fn.posterior_name)(mcmc_Y_dim, device, mcmc_train_Y_mean, mcmc_train_Y_std) bnn_post.set_sliced_pred(torch.tensor(mcmc_pred)) n_walkers = test_cfg.numerics.mcmc.walkerRatio * ( mcmc_Y_dim + 1) # BNN params + H0 times walker ratio init_pos = bnn_post.sample( n_walkers, sample_seed=test_cfg.global_seed ) # [batch_size, n_walkers, mcmc_Y_dim] contains just the lens model params, no D_dt init_D_dt = np.random.uniform(0.0, 10000.0, size=(batch_size, n_walkers, 1)) # FIXME: init H0 hardcoded kwargs_model = dict(lens_model_list=['PEMD', 'SHEAR'], point_source_model_list=['SOURCE_POSITION'], source_light_model_list=['SERSIC_ELLIPSE']) astro_sig = test_cfg.image_position_likelihood.sigma # Get H0 samples for each system if not test_cfg.time_delay_likelihood.baobab_time_delays: if 'abcd_ordering_i' not in master_truth: raise ValueError( "If the time delay measurements were not generated using Baobab, the user must specify the order of image positions in which the time delays are listed, in order of increasing dec." ) lenses_skipped = [] # keeps track of lenses that skipped MCMC total_progress = tqdm(total=n_test) # For each lens system... for i, lens_i in enumerate(lens_range): # Each lens gets a unique random state for td and vd measurement error realizations. rs_lens = np.random.RandomState(lens_i) ########################### # Relevant data and prior # ########################### data_i = master_truth.iloc[lens_i].copy() parameter_penalty.set_bnn_post_params( mcmc_pred[lens_i, :]) # set the BNN parameters # Init values for the lens model params if test_cfg.lens_posterior_type == 'default': init_info = dict( zip( mcmc_Y_cols, mcmc_pred[lens_i, :len(mcmc_Y_cols)] * mcmc_train_Y_std + mcmc_train_Y_mean)) # mean of primary Gaussian else: # types 'hybrid_with_truth_mean' and 'truth' init_info = dict(zip(mcmc_Y_cols, data_i[mcmc_Y_cols].values)) # truth params if not test_cfg.h0_posterior.exclude_velocity_dispersion: parameter_penalty.set_vel_disp_params() raise NotImplementedError lcdm = LCDM(z_lens=data_i['z_lens'], z_source=data_i['z_src'], flat=True) # Data is BCD - A with a certain ABCD ordering, so inferred time delays should follow this convention. measured_td_wrt0 = np.array(data_i['measured_td']) # [n_img - 1,] measured_td_sig = np.array(data_i['measured_td_err']) # [n_img - 1,] abcd_ordering_i = np.array(data_i['abcd_ordering_i']) n_img = len(abcd_ordering_i) kwargs_data_joint = dict( time_delays_measured=measured_td_wrt0, time_delays_uncertainties=measured_td_sig, abcd_ordering_i=abcd_ordering_i, #vel_disp_measured=measured_vd, # TODO: optionally exclude #vel_disp_uncertainty=vel_disp_sig, ) if not test_cfg.h0_posterior.exclude_velocity_dispersion: measured_vd = data_i['true_vd'] * ( 1.0 + rs_lens.randn() * test_cfg.error_model.velocity_dispersion_frac_error) kwargs_data_joint['vel_disp_measured'] = measured_vd kwargs_data_joint[ 'vel_disp_sig'] = test_cfg.velocity_dispersion_likelihood.sigma ############################# # Parameter init and bounds # ############################# lens_kwargs = mcmc_utils.get_lens_kwargs(init_info, null_spread=null_spread) ps_kwargs = mcmc_utils.get_ps_kwargs_src_plane(init_info, astro_sig, null_spread=null_spread) src_light_kwargs = mcmc_utils.get_light_kwargs( init_info['src_light_R_sersic'], null_spread=null_spread) special_kwargs = mcmc_utils.get_special_kwargs( n_img, astro_sig, null_spread=null_spread ) # image position offset and time delay distance, aka the "special" parameters kwargs_params = { 'lens_model': lens_kwargs, 'point_source_model': ps_kwargs, 'source_model': src_light_kwargs, 'special': special_kwargs, } if test_cfg.numerics.solver_type == 'NONE': solver_type = 'NONE' else: solver_type = 'PROFILE_SHEAR' if n_img == 4 else 'CENTER' #solver_type = 'NONE' kwargs_constraints = { 'num_point_source_list': [n_img], 'Ddt_sampling': True, 'solver_type': solver_type, } kwargs_likelihood = { 'time_delay_likelihood': True, 'sort_images_by_dec': True, 'prior_lens': [], 'prior_special': [], 'check_bounds': True, 'check_matched_source_position': False, 'source_position_tolerance': 0.01, 'source_position_sigma': 0.01, 'source_position_likelihood': False, 'custom_logL_addition': custom_logL_addition, } ########################### # MCMC posterior sampling # ########################### fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, verbose=False, mpi=False) if i == 0: param_class = fitting_seq._updateManager.param_class n_params, param_class_Y_cols = param_class.num_param() init_pos = mcmc_utils.reorder_to_param_class( mcmc_Y_cols, param_class_Y_cols, init_pos, init_D_dt) # MCMC sample from the post-processed BNN posterior jointly with cosmology lens_i_start_time = time.time() if test_cfg.lens_posterior_type == 'default': test_cfg.numerics.mcmc.update(init_samples=init_pos[lens_i, :, :]) fitting_kwargs_list_mcmc = [['MCMC', test_cfg.numerics.mcmc]] #with HiddenPrints(): try: chain_list_mcmc = fitting_seq.fit_sequence( fitting_kwargs_list_mcmc) kwargs_result_mcmc = fitting_seq.best_fit() except: print("lens {:d} skipped".format(lens_i)) total_progress.update(1) lenses_skipped.append(lens_i) continue lens_i_end_time = time.time() inference_time = (lens_i_end_time - lens_i_start_time) / 60.0 # min ############################# # Plotting the MCMC samples # ############################# # sampler_type : 'EMCEE' # samples_mcmc : np.array of shape `[n_mcmc_eval, n_params]` # param_mcmc : list of str of length n_params, the parameter names sampler_type, samples_mcmc, param_mcmc, _ = chain_list_mcmc[0] new_samples_mcmc = mcmc_utils.postprocess_mcmc_chain( kwargs_result_mcmc, samples_mcmc, kwargs_model, lens_kwargs[2], ps_kwargs[2], src_light_kwargs[2], special_kwargs[2], kwargs_constraints) # Plot D_dt histogram D_dt_samples = new_samples_mcmc['D_dt'].values true_D_dt = lcdm.D_dt(H_0=data_i['H0'], Om0=0.27) data_i['D_dt'] = true_D_dt # Export D_dt samples for this lens lens_inference_dict = dict( D_dt_samples=D_dt_samples, # kappa_ext=0 for these samples inference_time=inference_time, true_D_dt=true_D_dt, ) lens_inference_dict_save_path = os.path.join( out_dir, 'D_dt_dict_{0:04d}.npy'.format(lens_i)) np.save(lens_inference_dict_save_path, lens_inference_dict) # Optionally export the MCMC samples if test_cfg.export.mcmc_samples: mcmc_samples_path = os.path.join( out_dir, 'mcmc_samples_{0:04d}.csv'.format(lens_i)) new_samples_mcmc.to_csv(mcmc_samples_path, index=None) # Optionally export the D_dt histogram if test_cfg.export.D_dt_histogram: cleaned_D_dt_samples = h0_utils.remove_outliers_from_lognormal( D_dt_samples, 3) _ = plotting_utils.plot_D_dt_histogram(cleaned_D_dt_samples, lens_i, true_D_dt, save_dir=out_dir) # Optionally export the plot of MCMC chain if test_cfg.export.mcmc_chain: mcmc_chain_path = os.path.join( out_dir, 'mcmc_chain_{0:04d}.png'.format(lens_i)) plotting_utils.plot_mcmc_chain(chain_list_mcmc, mcmc_chain_path) # Optionally export posterior cornerplot of select lens model parameters with D_dt if test_cfg.export.mcmc_corner: mcmc_corner_path = os.path.join( out_dir, 'mcmc_corner_{0:04d}.png'.format(lens_i)) plotting_utils.plot_mcmc_corner( new_samples_mcmc[test_cfg.export.mcmc_cols], None, test_cfg.export.mcmc_col_labels, mcmc_corner_path) total_progress.update(1) total_progress.close()
} kwargs_likelihood = {'check_bounds': True, 'force_no_add_image': False, 'source_marg': False, 'image_position_uncertainty': 0.004, 'check_matched_source_position': True, 'source_position_tolerance': 0.001, 'time_delay_likelihood': True, } kwargs_numerics = {'supersampling_factor': 1} 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', 'time_delays_measured': delta_t["delta_t"].to_numpy(), 'time_delays_uncertainties': delta_t["sigma"].to_numpy(),} from lenstronomy.Workflow.fitting_sequence import FittingSequence fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) fitting_kwargs_list = [ ['PSO', {'sigma_scale': .1, 'n_particles': 200, 'n_iterations': 200}], ['MCMC', {'n_burn': 100, 'n_run': 100, 'walkerRatio': 10, 'sigma_scale': .1}] ] start_time = time.time() chain_list = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_seq.best_fit() end_time = time.time() print(end_time - start_time, 'total time needed for computation') print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ')
def test_minimizer(self): n_p = 2 n_i = 2 fitting_list = [] kwargs_simplex = {'n_iterations': n_i, 'method': 'Nelder-Mead'} fitting_list.append(['SIMPLEX', kwargs_simplex]) kwargs_simplex = {'n_iterations': n_i, 'method': 'Powell'} fitting_list.append(['SIMPLEX', kwargs_simplex]) kwargs_pso = { 'sigma_scale': 1, 'n_particles': n_p, 'n_iterations': n_i } fitting_list.append(['PSO', kwargs_pso]) kwargs_mcmc = { 'sigma_scale': 1, 'n_burn': 1, 'n_run': 1, 'n_walkers': 10, 'sampler_type': 'EMCEE' } fitting_list.append(['MCMC', kwargs_mcmc]) kwargs_mcmc['re_use_samples'] = True kwargs_mcmc['init_samples'] = np.array([[np.random.normal(1, 0.001)] for i in range(100)]) fitting_list.append(['MCMC', kwargs_mcmc]) def custom_likelihood(kwargs_lens, kwargs_source=None, kwargs_lens_light=None, kwargs_ps=None, kwargs_special=None, kwargs_extinction=None): theta_E = kwargs_lens[0]['theta_E'] return -(theta_E - 1.)**2 / 0.1**2 / 2 kwargs_likelihood = {'custom_logL_addition': custom_likelihood} kwargs_data_joint = {'multi_band_list': []} kwargs_model = {'lens_model_list': ['SIS']} kwargs_constraints = {} lens_param = [{ 'theta_E': 1, 'center_x': 0, 'center_y': 0 }], [{ 'theta_E': 0.1, 'center_x': 0.1, 'center_y': 0.1 }], [{ 'center_x': 0, 'center_y': 0 }], [{ 'theta_E': 0, 'center_x': -10, 'center_y': -10 }], [{ 'theta_E': 10, 'center_x': 10, 'center_y': 10 }] kwargs_params = {'lens_model': lens_param} fittingSequence = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) args = fittingSequence.param_class.kwargs2args( kwargs_lens=[{ 'theta_E': 1, 'center_x': 0, 'center_y': 0 }]) kwargs_result = fittingSequence.param_class.args2kwargs(args) print(kwargs_result) print(args, 'test args') chain_list = fittingSequence.fit_sequence(fitting_list) kwargs_result = fittingSequence.best_fit(bijective=False) npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'], 1, decimal=2)
def sl_sys_analysis(): # Get command line arguments args = {} if comm_rank == 0: print(":Registered %d processes" % comm_size) args["infile"] = sys.argv[1] args["nimgs"] = sys.argv[2] args["los"] = sys.argv[3] args["version"] = sys.argv[4] args["dt_sigma"] = float(sys.argv[5]) args["image_amps_sigma"] = float(sys.argv[6]) args["flux_ratio_errors"] = float(sys.argv[7]) args["astrometry_sigma"] = float(sys.argv[8]) args = comm.bcast(args) # Organize devision of strong lensing systems with open(args["infile"], "r") as myfile: limg_data = myfile.read() systems = json.loads(limg_data) sys_nr_per_proc = int(len(systems) / comm_size) print('comm_rank', comm_rank) start_sys = sys_nr_per_proc * comm_rank end_sys = sys_nr_per_proc * (comm_rank + 1) print(start_sys, end_sys) with open("../lens_catalogs_sie_only.json", "r") as myfile: limg_data = myfile.read() systems_prior = json.loads(limg_data) if comm_rank == 0: print("Each process will have %d systems" % sys_nr_per_proc) print("That should take app. %f min." % (sys_nr_per_proc * 20)) results = { "gamma": [], "phi_ext": [], "gamma_ext": [], "theta_E": [], "D_dt": [], } for ii in range(len(systems))[(start_sys + 2):end_sys]: system = systems[ii] system_prior = systems_prior[ii] print("Analysing system ID: %d" % ii) print(system) # the data set is z_lens = system_prior["zl"] z_source = 2.0 # multiple images properties ximg = np.zeros(system["nimgs"]) yimg = np.zeros(system["nimgs"]) delay = np.zeros(system["nimgs"]) image_amps = np.zeros(system["nimgs"]) for jj in range(system["nimgs"]): ximg[jj] = system["ximg"][jj] #[arcsec] yimg[jj] = system["yimg"][jj] #[arcsec] delay[jj] = system["delay"][jj] #[days] image_amps[jj] = system["mags"][jj] #[linear units or magnitudes] # sort by arrival time index_sort = np.argsort(delay) ximg = ximg[index_sort] yimg = yimg[index_sort] delay = delay[index_sort] image_amps = image_amps[index_sort] d_dt = delay[1:] - delay[0] # measurement uncertainties d_dt_sigma = np.ones(system["nimgs"] - 1) * args["dt_sigma"] image_amps_sigma = np.ones(system["nimgs"]) * args["image_amps_sigma"] flux_ratios = image_amps[1:] - image_amps[0] flux_ratio_errors = np.ones(system["nimgs"] - 1) * args["flux_ratio_errors"] # lens model choices lens_model_list = ["SPEP", "SHEAR"] # first choice: SPEP fixed_lens = [] kwargs_lens_init = [] kwargs_lens_sigma = [] kwargs_lower_lens = [] kwargs_upper_lens = [] fixed_lens.append({}) kwargs_lens_init.append({ "theta_E": 1.0, "gamma": 2, "center_x": 0, "center_y": 0, "e1": 0, "e2": 0.0, }) # error kwargs_lens_sigma.append({ "theta_E": 0.2, "e1": 0.1, "e2": 0.1, "gamma": 0.1, "center_x": 0.1, "center_y": 0.1, }) # lower limit kwargs_lower_lens.append({ "theta_E": 0.01, "e1": -0.5, "e2": -0.5, "gamma": 1.5, "center_x": -10, "center_y": -10, }) # upper limit kwargs_upper_lens.append({ "theta_E": 10, "e1": 0.5, "e2": 0.5, "gamma": 2.5, "center_x": 10, "center_y": 10, }) # second choice: SHEAR fixed_lens.append({"ra_0": 0, "dec_0": 0}) kwargs_lens_init.append({"e1": 0.0, "e2": 0.0}) kwargs_lens_sigma.append({"e1": 0.1, "e2": 0.1}) kwargs_lower_lens.append({"e1": -0.2, "e2": -0.2}) kwargs_upper_lens.append({"e1": 0.2, "e2": 0.2}) lens_params = [ kwargs_lens_init, kwargs_lens_sigma, fixed_lens, kwargs_lower_lens, kwargs_upper_lens, ] point_source_list = ["LENSED_POSITION"] fixed_ps = [{"ra_image": ximg, "dec_image": yimg}] kwargs_ps_init = fixed_ps # let some freedome in how well the actual image positions are # matching those given by the data (indicated as 'ra_image', 'dec_image' # and held fixed while fitting) kwargs_ps_sigma = [{ "ra_image": 0.01 * np.ones(len(ximg)), "dec_image": 0.01 * np.ones(len(ximg)), }] kwargs_lower_ps = [{ "ra_image": -10 * np.ones(len(ximg)), "dec_image": -10 * np.ones(len(ximg)), }] kwargs_upper_ps = [{ "ra_image": 10 * np.ones(len(ximg)), "dec_image": 10 * np.ones(len(ximg)) }] ps_params = [ kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps, ] fixed_cosmo = {} kwargs_cosmo_init = { "D_dt": 5000, "delta_x_image": np.zeros_like(ximg), "delta_y_image": np.zeros_like(ximg), } kwargs_cosmo_sigma = { "D_dt": 10000, "delta_x_image": np.ones_like(ximg) * args["astrometry_sigma"], "delta_y_image": np.ones_like(ximg) * args["astrometry_sigma"], } kwargs_lower_cosmo = { "D_dt": 0, "delta_x_image": np.ones_like(ximg) * (-1), "delta_y_image": np.ones_like(ximg) * (-1), } kwargs_upper_cosmo = { "D_dt": 10000, "delta_x_image": np.ones_like(ximg) * (1), "delta_y_image": np.ones_like(ximg) * (1), } cosmo_params = [ kwargs_cosmo_init, kwargs_cosmo_sigma, fixed_cosmo, kwargs_lower_cosmo, kwargs_upper_cosmo, ] kwargs_params = { "lens_model": lens_params, "point_source_model": ps_params, "cosmography": cosmo_params, } # setup options for likelihood and parameter sampling kwargs_constraints = { "num_point_source_list": [int(args["nimgs"])], # any proposed lens model must satisfy the image positions # appearing at the position of the point sources being sampeld "solver_type": "PROFILE_SHEAR", "cosmo_type": "D_dt", # sampling of the time-delay distance # explicit modelling of the astrometric imperfection of # the point source positions "point_source_offset": True, } kwargs_likelihood = { "check_bounds": True, "point_source_likelihood": True, "position_uncertainty": args["astrometry_sigma"], "check_solver": True, "solver_tolerance": 0.001, "time_delay_likelihood": True, "image_likelihood": False, # this needs to be explicitly given when not having imaging data "flux_ratio_likelihood": True, # enables the flux ratio likelihood } kwargs_data_joint = { "time_delays_measured": d_dt, "time_delays_uncertainties": d_dt_sigma, "flux_ratios": flux_ratios, "flux_ratio_errors": flux_ratio_errors, } kwargs_model = { "lens_model_list": lens_model_list, "point_source_model_list": point_source_list, } mpi = False # MPI possible, but not supported through that notebook. fitting_seq = FittingSequence( kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, ) fitting_kwargs_list = [ # ['update_settings', {'lens_add_fixed': [[0, ['gamma']]]}], [ "PSO", { "sigma_scale": 1.0, "n_particles": 100, "n_iterations": 100 } ] ] 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_result = ( fitting_seq.best_fit()) # and now we run the MCMC fitting_kwargs_list = [ [ "PSO", { "sigma_scale": 0.1, "n_particles": 100, "n_iterations": 100 } ], [ "MCMC", { "n_burn": 200, "n_run": 200, "walkerRatio": 10, "sigma_scale": 0.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_result = ( fitting_seq.best_fit()) # print("number of non-linear parameters in the MCMC process: ", len(param_mcmc)) # print("parameters in order: ", param_mcmc) print("number of evaluations in the MCMC process: ", np.shape(samples_mcmc)[0]) param = Param( kwargs_model, fixed_lens, kwargs_fixed_ps=fixed_ps, kwargs_fixed_cosmo=fixed_cosmo, kwargs_lens_init=lens_result, **kwargs_constraints, ) # the number of non-linear parameters and their names num_param, param_list = param.num_param() lensAnalysis = LensAnalysis(kwargs_model) mcmc_new_list = [] labels_new = [ r"$\gamma$", r"$\phi_{ext}$", r"$\gamma_{ext}$", r"$D_{\Delta t}$", ] D_dt = np.zeros(len(samples_mcmc)) theta_E = np.zeros(len(samples_mcmc)) for i in range(len(samples_mcmc)): # transform the parameter position of the MCMC chain in a # lenstronomy convention with keyword arguments kwargs_lens_out, kwargs_light_source_out, kwargs_light_lens_out, kwargs_ps_out, kwargs_cosmo = param.args2kwargs( samples_mcmc[i]) D_dt[i] = kwargs_cosmo["D_dt"] gamma[i] = kwargs_lens_out[0]["gamma"] theta_E[i] = kwargs_lens_out[0]['theta_E'] e1[i] = kwargs_lens_out[0]['e1'] e2[i] = kwargs_lens_out[0]['e2'] phi_ext, gamma_ext = lensAnalysis._lensModelExtensions.external_shear( kwargs_lens_out) # plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True) # and here the predicted angular diameter distance from a # default cosmology (attention for experimenter bias!) cosmo = FlatLambdaCDM( H0=71, Om0=0.3089, Ob0=0., ) lensCosmo = LensCosmo( cosmo=cosmo, z_lens=z_lens, z_source=z_source, ) gamma = np.mean(gamma) phi_ext = np.mean(phi_ext) gamma_ext = np.mean(gamma_ext) theta_E = np.mean(theta_E) D_dt = np.mean(D_dt) results["gamma"].append(gamma) results["phi_ext"].append(phi_ext) results["gamma_ext"].append(gamma_ext) results["theta_E"].append(theta_E) results["D_dt"].append(lensCosmo.D_dt) with open( "./quasars_%s_nimgs_%s_%s.json" % (args["los"], args["nimgs"], args["version"]), 'w') as fout: json.dump(results, fout)
def sl_sys_analysis(): # Get command line arguments args = {} if comm_rank == 0: print(":Registered %d processes" % comm_size) args["infile"] = sys.argv[1] args["nimgs"] = sys.argv[2] args["los"] = sys.argv[3] args["version"] = sys.argv[4] args["dt_sigma"] = float(sys.argv[5]) args["image_amps_sigma"] = float(sys.argv[6]) args["flux_ratio_errors"] = float(sys.argv[7]) args["astrometry_sigma"] = float(sys.argv[8]) args = comm.bcast(args) # Organize devision of strong lensing systems with open(args["infile"], "r") as myfile: limg_data = myfile.read() systems = json.loads(limg_data) sys_nr_per_proc = int(len(systems) / comm_size) print("comm_rank", comm_rank) start_sys = sys_nr_per_proc * comm_rank end_sys = sys_nr_per_proc * (comm_rank + 1) print(start_sys, end_sys) with open("../lens_catalogs_sie_only.json", "r") as myfile: limg_data = myfile.read() systems_prior = json.loads(limg_data) if comm_rank == 0: print("Each process will have %d systems" % sys_nr_per_proc) print("That should take app. %f min." % (sys_nr_per_proc * 20)) source_size_pc = 10.0 window_size = 0.1 # units of arcseconds grid_number = 100 # supersampled window (per axis) z_source = 2.0 cosmo = FlatLambdaCDM(H0=71, Om0=0.3089, Ob0=0.0) results = {"gamma": [], "phi_ext": [], "gamma_ext": [], "theta_E": [], "D_dt": []} for ii in range(len(systems))[(start_sys + 2) : end_sys]: system = systems[ii] system_prior = systems_prior[ii] print("Analysing system ID: %d" % ii) # the data set is z_lens = system_prior["zl"] lensCosmo = LensCosmo(cosmo=cosmo, z_lens=z_lens, z_source=z_source) # convert units of pc into arcseconds D_s = lensCosmo.D_s source_size_arcsec = source_size_pc / 10 ** 6 / D_s / constants.arcsec print("The source size in arcsec init = %.4f" % source_size_arcsec) #0.0012 # multiple images properties ximg = np.zeros(system["nimgs"]) yimg = np.zeros(system["nimgs"]) t_days = np.zeros(system["nimgs"]) image_amps = np.zeros(system["nimgs"]) for jj in range(system["nimgs"]): ximg[jj] = system["ximg"][jj] # [arcsec] yimg[jj] = system["yimg"][jj] # [arcsec] t_days[jj] = system["delay"][jj] # [days] image_amps[jj] = system["mags"][jj] # [linear units or magnitudes] # sort by arrival time index_sort = np.argsort(t_days) ximg = ximg[index_sort] # relative RA (arc seconds) yimg = yimg[index_sort] # relative DEC (arc seconds) image_amps = np.abs(image_amps[index_sort]) t_days = t_days[index_sort] d_dt = t_days[1:] - t_days[0] # measurement uncertainties astrometry_sigma = args["astrometry_sigma"] ximg_measured = ximg + np.random.normal(0, astrometry_sigma, system["nimgs"]) yimg_measured = yimg + np.random.normal(0, astrometry_sigma, system["nimgs"]) image_amps_sigma = np.ones(system["nimgs"]) * args["image_amps_sigma"] flux_ratios = image_amps[1:] - image_amps[0] flux_ratio_errors = np.ones(system["nimgs"] - 1) * args["flux_ratio_errors"] flux_ratios_measured = flux_ratios + np.random.normal(0, flux_ratio_errors) d_dt_sigma = np.ones(system["nimgs"] - 1) * args["dt_sigma"] d_dt_measured = d_dt + np.random.normal(0, d_dt_sigma) kwargs_data_joint = { "time_delays_measured": d_dt_measured, "time_delays_uncertainties": d_dt_sigma, "flux_ratios": flux_ratios_measured, "flux_ratio_errors": flux_ratio_errors, "ra_image_list": [ximg_measured], "dec_image_list": [yimg_measured], } # lens model choices lens_model_list = ["SPEMD", "SHEAR_GAMMA_PSI"] # 1. layer: primary SPEP fixed_lens = [] kwargs_lens_init = [] kwargs_lens_sigma = [] kwargs_lower_lens = [] kwargs_upper_lens = [] fixed_lens.append({}) kwargs_lens_init.append( { "theta_E": 1.0, "gamma": 2, "center_x": 0, "center_y": 0, "e1": 0, "e2": 0.0, } ) # error kwargs_lens_sigma.append( { "theta_E": 0.2, "e1": 0.1, "e2": 0.1, "gamma": 0.1, "center_x": 0.1, "center_y": 0.1, } ) # lower limit kwargs_lower_lens.append( { "theta_E": 0.01, "e1": -0.5, "e2": -0.5, "gamma": 1.5, "center_x": -10, "center_y": -10, } ) # upper limit kwargs_upper_lens.append( { "theta_E": 10, "e1": 0.5, "e2": 0.5, "gamma": 2.5, "center_x": 10, "center_y": 10, } ) # 2nd layer: external SHEAR fixed_lens.append({"ra_0": 0, "dec_0": 0}) kwargs_lens_init.append({"gamma_ext": 0.05, "psi_ext": 0.0}) kwargs_lens_sigma.append({"gamma_ext": 0.05, "psi_ext": np.pi}) kwargs_lower_lens.append({"gamma_ext": 0, "psi_ext": -np.pi}) kwargs_upper_lens.append({"gamma_ext": 0.3, "psi_ext": np.pi}) # 3rd layer: external CONVERGENCE kwargs_lens_init.append({'kappa_ext': 0.12}) kwargs_lens_sigma.append({'kappa_ext': 0.06}) kwargs_lower_lens.append({'kappa_ext': 0.0}) kwargs_upper_lens.append({'kappa_ext': 0.3}) # combined lens model lens_params = [ kwargs_lens_init, kwargs_lens_sigma, fixed_lens, kwargs_lower_lens, kwargs_upper_lens, ] # image position parameters point_source_list = ["LENSED_POSITION"] # we fix the image position coordinates fixed_ps = [{}] # the initial guess for the appearing image positions is: # at the image position. kwargs_ps_init = [{"ra_image": ximg, "dec_image": yimg}] # let some freedome in how well the actual image positions are # matching those given by the data (indicated as 'ra_image', 'dec_image' # and held fixed while fitting) kwargs_ps_sigma = [ { "ra_image": 0.01 * np.ones(len(ximg)), "dec_image": 0.01 * np.ones(len(ximg)), } ] kwargs_lower_ps = [ { "ra_image": -10 * np.ones(len(ximg)), "dec_image": -10 * np.ones(len(ximg)), } ] kwargs_upper_ps = [ {"ra_image": 10 * np.ones(len(ximg)), "dec_image": 10 * np.ones(len(ximg))} ] ps_params = [ kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps, ] # quasar source size fixed_special = {} kwargs_special_init = {} kwargs_special_sigma = {} kwargs_lower_special = {} kwargs_upper_special = {} fixed_special["source_size"] = source_size_arcsec kwargs_special_init["source_size"] = source_size_arcsec kwargs_special_sigma["source_size"] = source_size_arcsec kwargs_lower_special["source_size"] = 0.0001 kwargs_upper_special["source_size"] = 1 # Time-delay distance kwargs_special_init["D_dt"] = 4300 # corresponds to H0 ~ 70 kwargs_special_sigma["D_dt"] = 3000 kwargs_lower_special["D_dt"] = 2500 # corresponds to H0 ~ 120 kwargs_upper_special["D_dt"] = 14000 # corresponds to H0 ~ 20 special_params = [ kwargs_special_init, kwargs_special_sigma, fixed_special, kwargs_lower_special, kwargs_upper_special, ] # combined parameter settings kwargs_params = { "lens_model": lens_params, "point_source_model": ps_params, "special": special_params, } # our model choices kwargs_model = { "lens_model_list": lens_model_list, "point_source_model_list": point_source_list, } lensModel = LensModel(kwargs_model["lens_model_list"]) lensModelExtensions = LensModelExtensions(lensModel=lensModel) lensEquationSolver = LensEquationSolver(lensModel=lensModel) # setup options for likelihood and parameter sampling time_delay_likelihood = True flux_ratio_likelihood = True image_position_likelihood = True kwargs_flux_compute = { "source_type": "INF", "window_size": window_size, "grid_number": grid_number, } kwargs_constraints = { "num_point_source_list": [int(args["nimgs"])], # any proposed lens model must satisfy the image positions # appearing at the position of the point sources being sampeld # "solver_type": "PROFILE_SHEAR", "Ddt_sampling": time_delay_likelihood, # sampling of the time-delay distance # explicit modelling of the astrometric imperfection of # the point source positions "point_source_offset": True, } # explicit sampling of finite source size parameter # (only use when source_type='GAUSSIAN' or 'TORUS') if ( kwargs_flux_compute["source_type"] in ["GAUSSIAN", "TORUS"] and flux_ratio_likelihood is True ): kwargs_constraints["source_size"] = True # e.g. power-law mass slope of the main deflector # [[index_model, 'param_name', mean, 1-sigma error], [...], ...] prior_lens = [[0, "gamma", 2, 0.1]] prior_special = [] kwargs_likelihood = { "position_uncertainty": args["astrometry_sigma"], "source_position_likelihood": True, "image_position_likelihood": True, "time_delay_likelihood": True, "flux_ratio_likelihood": True, "kwargs_flux_compute": kwargs_flux_compute, "prior_lens": prior_lens, "prior_special": prior_special, "check_solver": True, "solver_tolerance": 0.001, "check_bounds": True, } fitting_seq = FittingSequence( kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, ) fitting_kwargs_list = [ ["PSO", {"sigma_scale": 1.0, "n_particles": 200, "n_iterations": 500}] ] chain_list_pso = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_seq.best_fit() kwargs_result = fitting_seq.best_fit(bijective=True) args_result = fitting_seq.param_class.kwargs2args(**kwargs_result) logL, _ = fitting_seq.likelihoodModule.logL(args_result, verbose=True) # and now we run the MCMC fitting_kwargs_list = [ [ "MCMC", {"n_burn": 400, "n_run": 600, "walkerRatio": 10, "sigma_scale": 0.1}, ] ] chain_list_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_seq.best_fit() # print("number of non-linear parameters in the MCMC process: ", len(param_mcmc)) # print("parameters in order: ", param_mcmc) print("number of evaluations in the MCMC process: ", np.shape(samples_mcmc)[0]) param = Param( kwargs_model, fixed_lens, kwargs_fixed_ps=fixed_ps, kwargs_fixed_special=fixed_special, kwargs_lens_init=kwargs_result["kwargs_lens"], **kwargs_constraints, ) # the number of non-linear parameters and their names # num_param, param_list = param.num_param() for i in range(len(samples_mcmc)): kwargs_out = param.args2kwargs(samples_mcmc[i]) kwargs_lens_out, kwargs_special_out, kwargs_ps_out = ( kwargs_out["kwargs_lens"], kwargs_out["kwargs_special"], kwargs_out["kwargs_ps"], ) # compute 'real' image position adding potential astrometric shifts x_pos = kwargs_ps_out[0]["ra_image"] y_pos = kwargs_ps_out[0]["dec_image"] # extract quantities of the main deflector theta_E = kwargs_lens_out[0]["theta_E"] gamma = kwargs_lens_out[0]["gamma"] e1, e2 = kwargs_lens_out[0]["e1"], kwargs_lens_out[0]["e2"] phi, q = param_util.ellipticity2phi_q(e1, e2) phi_ext, gamma_ext = ( kwargs_lens_out[1]["psi_ext"] % np.pi, kwargs_lens_out[1]["gamma_ext"], ) if flux_ratio_likelihood is True: mag = lensModel.magnification(x_pos, y_pos, kwargs_lens_out) flux_ratio_fit = mag[1:] / mag[0] if ( kwargs_constraints.get("source_size", False) is True and "source_size" not in fixed_special ): source_size = kwargs_special_out["source_size"] if time_delay_likelihood is True: D_dt = kwargs_special_out["D_dt"] # and here the predicted angular diameter distance from a # default cosmology (attention for experimenter bias!) gamma = np.median(gamma) phi_ext = np.median(phi_ext) gamma_ext = np.median(gamma_ext) theta_E = np.median(theta_E) D_dt = np.median(D_dt) results["gamma"].append(gamma) results["phi_ext"].append(phi_ext) results["gamma_ext"].append(gamma_ext) results["theta_E"].append(theta_E) results["H0"].append(c_light / D_dt)
class ClsrWorkflow(object): def __init__(self, kwargs_data_joint, kwargs_model,lens_params,source_params, lenslight_params=None, kwargs_constraints=None, kwargs_likelihood=None): """ class to manage cluster source reconstruction. This class inherited the FittingSequence class in Workflow module of lenstronomy. :param kwargs_data_joint: keywords arguments of [data, psf, numericals] in lenstronomy convention. :param kwargs_model: name of model list :param lens_params: lens model keywords arguments [kwargs_lens_init, kwargs_lens_sigma, kwargs_fixed_lens, kwargs_lower_lens, kwargs_upper_lens] :param source_params: source model keywords arguments [kwargs_source_init, kwargs_source_sigma, kwargs_fixed_source, kwargs_lower_source, kwargs_upper_source] :param kwargs_constraints: contraints on models :param kwargs_likelihood: options of calculating likelihood, see more: LikelihoodModule class in Sampling module of lenstronomy. """ self.kwargs_data_joint =kwargs_data_joint self.multi_band_list = kwargs_data_joint.get('multi_band_list', []) self.kwargs_model =kwargs_model kwargs_params = {'lens_model': lens_params, 'source_model': source_params, 'lens_light_model': lenslight_params} self.kwargs_params= kwargs_params if kwargs_constraints is None: kwargs_constraints ={} if kwargs_likelihood is None: kwargs_likelihood = {'source_marg': False, 'check_positive_flux': True} self.fitting_seq_src = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) def run_fit_sequence(self,fitting_kwargs_list): """ :param fitting_kwargs_list: list of [['string', {kwargs}], ..] with 'string being the specific fitting option and kwargs being the arguments passed to this option :return: fitting results """ chain_list = self.fitting_seq_src.fit_sequence(fitting_kwargs_list) kwargs_result = self.fitting_seq_src.best_fit(bijective=False) bic_model = self.fitting_seq_src.bic return bic_model,chain_list, kwargs_result def lensmodel_comp(self, num_img, n_particles,n_iterations,sigma_scale, num_lens_model,fixed_index =0,flexion_option=True): """ function to figure out the necessary of increasing lens model complexity. Currently, we only consider up to flexion term. :param n_particles: particles numbers of PSO :param n_iterations: iteration numbers of PSO :param sigma_scale: sigma scale for PSO :param num_img: int, numbers of lensed image :param fixed_index: int, index of fixed lensed image :param num_lens_model: numbers of strings contained in lens model list :param flexion_option: bool, default is taking flexion into consideration :return: pso results, fitting results of necessary lens model complexity, bic values """ lens_remove_fixed_list = [] lens_add_fixed_list = [] for i in range(num_img): if i == fixed_index: print ("lens model keep fixed in frame:", i+1) else: lens_flexion_index = (i + 1) * num_lens_model - 1 lens_remove_fixed_list.append([lens_flexion_index, ['G1', 'G2', 'F1', 'F2'], [0, 0, 0, 0]]) G1_fix = self.kwargs_params['lens_model'][2][lens_flexion_index]['G1'] G2_fix = self.kwargs_params['lens_model'][2][lens_flexion_index]['G2'] F1_fix = self.kwargs_params['lens_model'][2][lens_flexion_index]['F1'] F2_fix = self.kwargs_params['lens_model'][2][lens_flexion_index]['F2'] lens_add_fixed_list.append([lens_flexion_index, ['G1', 'G2', 'F1', 'F2'], [G1_fix, G2_fix, F1_fix, F2_fix]]) flexion_add_fixed = [['update_settings', {'lens_add_fixed': lens_add_fixed_list}]] print("flexion_fixed:", flexion_add_fixed) kwargs_pso = [['PSO', {'sigma_scale': sigma_scale, 'n_particles': n_particles, 'n_iterations': n_iterations}]] fitting_kwargs_fix = flexion_add_fixed+ kwargs_pso bic_model_fix, chain_list_fix, kwargs_result_fix = self.run_fit_sequence(fitting_kwargs_fix) if flexion_option: flexion_remove_fixed = [['update_settings', {'lens_remove_fixed': lens_remove_fixed_list}]] print ("flexion_remove_fixed:", flexion_remove_fixed) fitting_kwargs_free = flexion_remove_fixed + kwargs_pso bic_model_free, chain_list_free, kwargs_result_free = self.run_fit_sequence(fitting_kwargs_free) else: bic_model_free = 10000000 if bic_model_free > bic_model_fix: print ("No necessary to add flexion!") bic_list = [bic_model_fix] chain_list = [chain_list_fix] kwargs_result_list = [kwargs_result_fix] self._update_kwargs(kwargs_result_fix) _, _, _ = self.run_fit_sequence(flexion_add_fixed) elif bic_model_free < bic_model_fix: print("Flexion is needed!") bic_list = [bic_model_fix, bic_model_free] chain_list = [chain_list_fix] kwargs_result_list = [kwargs_result_fix] return chain_list, kwargs_result_list, bic_list def sourcemodel_comp(self, n_max_range=[0], sr = 0, n_particles=10, n_iterations=10,sigma_scale =1.0, bic_model_in = [100000], chain_list_in = [0], kwargs_results_in = [0], bic_option=True) : """ function to found the best fitting results among source models related to numbers of shapelets basis :param n_max_range: shapelets basis selection range :param sr: typical scale (") in source plane :param n_particles: particles numbers of PSO :param n_iterations: iteration numbers of PSO :param sigma_scale: sigma scale for PSO :param bic_model_in: BIC value of models before adding shapelets to source model :param chain_list_in: PSO chain results of models before adding shapelets to source model :param kwargs_results_in: fitting results of models before adding shapelets to source model :return: best-fit PSO results, best-fits results, PSO results for n_max_range, fitting results for n_max_range, bic values for all models """ bic_model_list = bic_model_in chain_list_list = chain_list_in kwargs_result_list = kwargs_results_in bic_in_len = len(bic_model_in) bic_run = True beta0 = sr kwargs_pso = [['PSO', {'sigma_scale': sigma_scale, 'n_particles': n_particles, 'n_iterations': n_iterations}]] for nmax in n_max_range: if nmax < 0: raise ValueError("nmax can not be negative!",nmax) else: if nmax == n_max_range[0]: start_kwargs_shapelet = [['update_settings', {'source_remove_fixed': [ [1, ['beta'], [beta0]] ]}]] else: start_kwargs_shapelet = [] beta_nmax = ((nmax + 1)) ** 0.5 * beta0 fit_kwargs_shapelet = [['update_settings', {'source_add_fixed': [[1, ['n_max'], [nmax]]], 'change_source_lower_limit': [[1, ['beta'], [beta_nmax]]] } ]] fitting_kwargs = start_kwargs_shapelet + fit_kwargs_shapelet + kwargs_pso if bic_run: print ("nmax",nmax,"fitting_kwargs",fitting_kwargs) bic_model,chain_list, kwargs_result = self.run_fit_sequence(fitting_kwargs) if bic_model > bic_model_list[-1]: if bic_option: bic_run = False if bic_model > bic_model_in[bic_in_len-1]: print ("bic_model_in",bic_model_in) print ("no necessary to add SHAPELETS !") fix_kwargs_shapelet=[['update_settings', {'source_add_fixed': [[1, ['beta'], [sr]]]}]] _, _, _ = self.run_fit_sequence(fix_kwargs_shapelet) elif not bic_option: chain_list_list.append(chain_list) kwargs_result_list.append(kwargs_result) bic_model_list.append(bic_model) print ("no necessary to increase model complexity!") elif bic_model < bic_model_list[-1]: chain_list_list.append(chain_list) kwargs_result_list.append(kwargs_result) bic_model_list.append(bic_model) print (bic_model, "currently is the lowest BIC value in bic_model_list=", bic_model_list) bic_sourcemodel = bic_model_list[bic_in_len:] if bic_sourcemodel ==[]: chain_list_lowest = chain_list_in[-1] kwargs_result_lowest = kwargs_results_in[-1] else: index_bic_minima = np.where(bic_model_list == np.min(bic_model_list))[0][0] chain_list_lowest = chain_list_list[index_bic_minima] kwargs_result_lowest = kwargs_result_list[index_bic_minima] return chain_list_lowest, kwargs_result_lowest, chain_list_list, kwargs_result_list, bic_model_list def _update_kwargs(self,kwargs_result): """ :param kwargs_result: fitting results of a specific state :return: go back to a specific state """ self.fitting_seq_src.update_state(kwargs_result)
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
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 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 test_fitting_sequence(self): fittingSequence = FittingSequence(self.kwargs_data_joint, self.kwargs_model, self.kwargs_constraints, self.kwargs_likelihood, self.kwargs_params) kwargs_result = fittingSequence.best_fit(bijective=False) lens_temp = kwargs_result['kwargs_lens'] npt.assert_almost_equal(lens_temp[0]['theta_E'], self.kwargs_lens[0]['theta_E'], decimal=2) logL = fittingSequence.best_fit_likelihood print(logL, 'test') #print(lens_temp, source_temp, lens_light_temp, ps_temp, cosmo_temp) assert logL < 0 bic = fittingSequence.bic assert bic > 0 #npt.assert_almost_equal(bic, 20000000220.29376, decimal=-4) #npt.assert_almost_equal(logL, -10000000061.792593, decimal=-4) n_p = 2 n_i = 2 fitting_list = [] kwargs_pso = { 'sigma_scale': 1, 'n_particles': n_p, 'n_iterations': n_i } fitting_list.append(['PSO', kwargs_pso]) kwargs_align = { 'lowerLimit': -0.1, 'upperLimit': 0.1, 'n_particles': 2, 'n_iterations': 2 } fitting_list.append(['align_images', kwargs_align]) kwargs_psf_iter = { 'num_iter': 2, 'psf_iter_factor': 0.5, 'stacking_method': 'mean', 'new_procedure': False } fitting_list.append(['psf_iteration', kwargs_psf_iter]) fitting_list.append(['restart', None]) fitting_list.append(['fix_not_computed', {'free_bands': [True]}]) n_sersic_overwrite = 4 kwargs_update = { 'lens_light_add_fixed': [[0, ['n_sersic'], [n_sersic_overwrite]]], 'lens_light_remove_fixed': [[0, ['center_x']]], 'change_source_lower_limit': [[0, ['n_sersic'], [0.1]]], 'change_source_upper_limit': [[0, ['n_sersic'], [10]]] } fitting_list.append(['update_settings', kwargs_update]) chain_list = fittingSequence.fit_sequence(fitting_list) lens_fixed, source_fixed, lens_light_fixed, ps_fixed, special_fixed, extinction_fixed = fittingSequence._updateManager.fixed_kwargs kwargs_result = fittingSequence.best_fit(bijective=False) npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'], self.kwargs_lens[0]['theta_E'], decimal=1) npt.assert_almost_equal( fittingSequence._updateManager._lens_light_fixed[0]['n_sersic'], n_sersic_overwrite, decimal=8) npt.assert_almost_equal(lens_light_fixed[0]['n_sersic'], 4, decimal=-1) assert fittingSequence._updateManager._lower_kwargs[1][0][ 'n_sersic'] == 0.1 assert fittingSequence._updateManager._upper_kwargs[1][0][ 'n_sersic'] == 10 # test 'set_param_value' fitting sequence fitting_list = [[ 'set_param_value', { 'lens': [[1, ['gamma1'], [0.013]]] } ], ['set_param_value', { 'lens_light': [[0, ['center_x'], [0.009]]] }], ['set_param_value', { 'source': [[0, ['n_sersic'], [2.993]]] }], ['set_param_value', { 'ps': [[0, ['ra_source'], [0.007]]] }]] fittingSequence.fit_sequence(fitting_list) kwargs_set = fittingSequence._updateManager.parameter_state assert kwargs_set['kwargs_lens'][1]['gamma1'] == 0.013 assert kwargs_set['kwargs_lens_light'][0]['center_x'] == 0.009 assert kwargs_set['kwargs_source'][0]['n_sersic'] == 2.993 assert kwargs_set['kwargs_ps'][0]['ra_source'] == 0.007
def test_fitting_sequence(self): # kwargs_init = [self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps] lens_sigma = [{ 'theta_E': 0.1, 'gamma': 0.1, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1 }, { 'e1': 0.1, 'e2': 0.1 }] lens_lower = [{ 'theta_E': 0., 'gamma': 1.5, 'center_x': -2, 'center_y': -2, 'e1': -0.4, 'e2': -0.4 }, { 'e1': -0.3, 'e2': -0.3 }] lens_upper = [{ 'theta_E': 10., 'gamma': 2.5, 'center_x': 2, 'center_y': 2, 'e1': 0.4, 'e2': 0.4 }, { 'e1': 0.3, 'e2': 0.3 }] source_sigma = [{ 'R_sersic': 0.05, 'n_sersic': 0.5, 'center_x': 0.1, 'center_y': 0.1, 'e1': 0.1, 'e2': 0.1 }] source_lower = [{ 'R_sersic': 0.01, 'n_sersic': 0.5, 'center_x': -2, 'center_y': -2, 'e1': -0.4, 'e2': -0.4 }] source_upper = [{ 'R_sersic': 10, 'n_sersic': 5.5, 'center_x': 2, 'center_y': 2, 'e1': 0.4, 'e2': 0.4 }] lens_light_sigma = [{ 'R_sersic': 0.05, 'n_sersic': 0.5, 'center_x': 0.1, 'center_y': 0.1 }] lens_light_lower = [{ 'R_sersic': 0.01, 'n_sersic': 0.5, 'center_x': -2, 'center_y': -2 }] lens_light_upper = [{ 'R_sersic': 10, 'n_sersic': 5.5, 'center_x': 2, 'center_y': 2 }] ps_sigma = [{'ra_source': 1, 'dec_source': 1, 'point_amp': 1}] lens_param = self.kwargs_lens, lens_sigma, [{}, { 'ra_0': 0, 'dec_0': 0 }], lens_lower, lens_upper source_param = self.kwargs_source, source_sigma, [ {} ], source_lower, source_upper lens_light_param = self.kwargs_lens_light, lens_light_sigma, [{ 'center_x': 0 }], lens_light_lower, lens_light_upper ps_param = self.kwargs_ps, ps_sigma, [{} ], self.kwargs_ps, self.kwargs_ps kwargs_params = { 'lens_model': lens_param, 'source_model': source_param, 'lens_light_model': lens_light_param, 'point_source_model': ps_param, # 'cosmography': cosmo_param } # kwargs_params = [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_init, kwargs_init] image_band = [self.kwargs_data, self.kwargs_psf, self.kwargs_numerics] multi_band_list = [image_band] kwargs_data_joint = { 'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear' } fittingSequence = FittingSequence(kwargs_data_joint, self.kwargs_model, self.kwargs_constraints, self.kwargs_likelihood, kwargs_params) kwargs_result = fittingSequence.best_fit(bijective=False) lens_temp = kwargs_result['kwargs_lens'] npt.assert_almost_equal(lens_temp[0]['theta_E'], self.kwargs_lens[0]['theta_E'], decimal=2) logL = fittingSequence.best_fit_likelihood print(logL, 'test') #print(lens_temp, source_temp, lens_light_temp, ps_temp, cosmo_temp) assert logL < 0 bic = fittingSequence.bic assert bic > 0 #npt.assert_almost_equal(bic, 20000000220.29376, decimal=-4) #npt.assert_almost_equal(logL, -10000000061.792593, decimal=-4) n_p = 2 n_i = 2 fitting_list = [] kwargs_pso = { 'sigma_scale': 1, 'n_particles': n_p, 'n_iterations': n_i } fitting_list.append(['PSO', kwargs_pso]) kwargs_mcmc = { 'sigma_scale': 0.1, 'n_burn': 1, 'n_run': 1, 'walkerRatio': 2 } fitting_list.append(['MCMC', kwargs_mcmc]) kwargs_mcmc['re_use_samples'] = True fitting_list.append(['MCMC', kwargs_mcmc]) kwargs_mcmc['sampler_type'] = 'EMCEE' fitting_list.append(['MCMC', kwargs_mcmc]) kwargs_align = { 'lowerLimit': -0.1, 'upperLimit': 0.1, 'n_particles': 2, 'n_iterations': 2 } fitting_list.append(['align_images', kwargs_align]) kwargs_psf_iter = { 'num_iter': 2, 'psf_iter_factor': 0.5, 'stacking_method': 'mean' } fitting_list.append(['psf_iteration', kwargs_psf_iter]) fitting_list.append(['restart', None]) fitting_list.append(['fix_not_computed', {'free_bands': [True]}]) n_sersic_overwrite = 4 kwargs_update = { 'lens_light_add_fixed': [[0, ['n_sersic'], [n_sersic_overwrite]]], 'lens_light_remove_fixed': [[0, ['center_x']]], 'change_source_lower_limit': [[0, ['n_sersic'], [0.1]]], 'change_source_upper_limit': [[0, ['n_sersic'], [10]]] } fitting_list.append(['update_settings', kwargs_update]) #kwargs_model = {}, kwargs_constraints = {}, kwargs_likelihood = {}, lens_add_fixed = [], #source_add_fixed = [], lens_light_add_fixed = [], ps_add_fixed = [], cosmo_add_fixed = [], lens_remove_fixed = [], #source_remove_fixed = [], lens_light_remove_fixed = [], ps_remove_fixed = [], cosmo_remove_fixed = [] chain_list = fittingSequence.fit_sequence(fitting_list) lens_fixed, source_fixed, lens_light_fixed, ps_fixed, special_fixed, extinction_fixed = fittingSequence._updateManager._fixed_kwargs kwargs_result = fittingSequence.best_fit(bijective=False) npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'], self.kwargs_lens[0]['theta_E'], decimal=1) npt.assert_almost_equal( fittingSequence._updateManager._lens_light_fixed[0]['n_sersic'], n_sersic_overwrite, decimal=8) npt.assert_almost_equal(lens_light_fixed[0]['n_sersic'], 4, decimal=-1) assert fittingSequence._updateManager._lower_kwargs[1][0][ 'n_sersic'] == 0.1 assert fittingSequence._updateManager._upper_kwargs[1][0][ 'n_sersic'] == 10 # Nested sampler tests # further decrease the parameter space for nested samplers to run faster fitting_list2 = [] kwargs_update2 = { 'ps_add_fixed': [[0, ['ra_source', 'dec_source'], [0, 0]]], 'lens_light_add_fixed': [[ 0, ['n_sersic', 'R_sersic', 'center_x', 'center_y'], [4, .1, 0, 0] ]], 'source_add_fixed': [[ 0, ['R_sersic', 'e1', 'e2', 'center_x', 'center_y'], [.6, .1, .1, 0, 0] ]], 'lens_add_fixed': [[ 0, ['gamma', 'theta_E', 'e1', 'e2', 'center_x', 'center_y'], [1.8, 1., .1, .1, 0, 0] ], [1, ['e1', 'e2'], [0.01, 0.01]]], 'change_source_lower_limit': [[0, ['n_sersic'], [2.9]]], 'change_source_upper_limit': [[0, ['n_sersic'], [3.1]]] } fitting_list2.append(['update_settings', kwargs_update2]) kwargs_multinest = { 'sampler_type': 'MULTINEST', 'kwargs_run': { 'n_live_points': 10, 'evidence_tolerance': 0.5, 'sampling_efficiency': 0.8, # 1 for posterior-only, 0 for evidence-only 'importance_nested_sampling': False, 'multimodal': True, 'const_efficiency_mode': False, # reduce sampling_efficiency to 5% when True }, 'remove_output_dir': True, } fitting_list2.append(['nested_sampling', kwargs_multinest]) kwargs_dynesty = { 'sampler_type': 'DYNESTY', 'kwargs_run': { 'dlogz_init': 0.01, 'nlive_init': 3, 'nlive_batch': 3, 'maxbatch': 1, }, } fitting_list2.append(['nested_sampling', kwargs_dynesty]) kwargs_dypolychord = { 'sampler_type': 'DYPOLYCHORD', 'kwargs_run': { 'ninit': 8, 'nlive_const': 10, #'seed_increment': 1, 'resume_dyn_run': False, #'init_step': 10, }, 'polychord_settings': { 'seed': 1, #'num_repeats': 20 }, 'dypolychord_dynamic_goal': 0.8, # 1 for posterior-only, 0 for evidence-only 'remove_output_dir': True, } fitting_list2.append(['nested_sampling', kwargs_dypolychord]) chain_list2 = fittingSequence.fit_sequence(fitting_list2) kwargs_fixed = fittingSequence._updateManager._fixed_kwargs npt.assert_almost_equal(kwargs_fixed[0][1]['e1'], 0.01, decimal=2) assert fittingSequence._updateManager._lower_kwargs[1][0][ 'n_sersic'] == 2.9 assert fittingSequence._updateManager._upper_kwargs[1][0][ 'n_sersic'] == 3.1 kwargs_test = {'kwargs_lens': 1} fittingSequence.update_state(kwargs_test) kwargs_out = fittingSequence.best_fit(bijective=True) assert kwargs_out['kwargs_lens'] == 1
from lenstronomy.Workflow.fitting_sequence import FittingSequence fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params) fitting_kwargs_list = [ #['update_settings', {'lens_add_fixed': [[0, ['gamma']]]}], ['PSO', { 'sigma_scale': 1., 'n_particles': 100, 'n_iterations': 100 }], ] start_time = time.time() 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_result = fitting_seq.best_fit( ) end_time = time.time() print(end_time - start_time, 'total time needed for computation') print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ') #and now we run the MCMC fitting_kwargs_list = [[ 'PSO', { 'sigma_scale': .1, 'n_particles': 100, 'n_iterations': 100 } ], [ 'MCMC', { 'n_burn': 200, 'n_run': 200,
def swim(self, lens_name, model_id, log=True, mpi=False, recipe_name='default', sampler='EMCEE', thread_count=1): """ Run models for a single lens. :param lens_name: lens name :type lens_name: `str` :param model_id: identifier for the model run :type model_id: `str` :param log: if `True`, all `print` statements will be logged :type log: `bool` :param mpi: MPI option :type mpi: `bool` :param recipe_name: recipe for pre-sampling optimization, supported ones now: 'default' and 'galaxy-galaxy' :type recipe_name: `str` :param sampler: 'EMCEE' or 'COSMOHAMMER', cosmohammer is kept for legacy :type sampler: `str` :param thread_count: number of threads if `multiprocess` is used :type thread_count: `int` :return: :rtype: """ pool = choose_pool(mpi=mpi) if log and pool.is_master(): log_file = open(self.file_system.get_log_file_path(lens_name, model_id), 'wt') sys.stdout = log_file config = self.get_lens_config(lens_name) recipe = Recipe(config, sampler=sampler, thread_count=thread_count) psf_supersampling_factor = config.get_psf_supersampled_factor() kwargs_data_joint = self.get_kwargs_data_joint( lens_name, psf_supersampled_factor=psf_supersampling_factor) fitting_sequence = FittingSequence( kwargs_data_joint, config.get_kwargs_model(), config.get_kwargs_constraints(), config.get_kwargs_likelihood(), config.get_kwargs_params(), mpi=mpi ) fitting_kwargs_list = recipe.get_recipe( kwargs_data_joint=kwargs_data_joint, recipe_name=recipe_name) fit_output = fitting_sequence.fit_sequence(fitting_kwargs_list) kwargs_result = fitting_sequence.best_fit(bijective=False) output = { 'settings': config.settings, 'kwargs_result': kwargs_result, 'fit_output': fit_output, } if pool.is_master(): self.file_system.save_output(lens_name, model_id, output) if log and pool.is_master(): log_file.close()
def main(): args = script_utils.parse_inference_args() test_cfg = TestConfig.from_file(args.test_config_file_path) baobab_cfg = BaobabConfig.from_file(test_cfg.data.test_baobab_cfg_path) cfg = TrainValConfig.from_file(test_cfg.train_val_config_file_path) # Set device and default data type device = torch.device(test_cfg.device_type) if device.type == 'cuda': torch.set_default_tensor_type('torch.cuda.' + cfg.data.float_type) else: torch.set_default_tensor_type('torch.' + cfg.data.float_type) script_utils.seed_everything(test_cfg.global_seed) ############ # Data I/O # ############ train_data = XYData( is_train=True, Y_cols=cfg.data.Y_cols, float_type=cfg.data.float_type, define_src_pos_wrt_lens=cfg.data.define_src_pos_wrt_lens, rescale_pixels=cfg.data.rescale_pixels, rescale_pixels_type=cfg.data.rescale_pixels_type, log_pixels=cfg.data.log_pixels, add_pixel_noise=cfg.data.add_pixel_noise, eff_exposure_time=cfg.data.eff_exposure_time, train_Y_mean=None, train_Y_std=None, train_baobab_cfg_path=cfg.data.train_baobab_cfg_path, val_baobab_cfg_path=None, for_cosmology=False) # Define val data and loader test_data = XYData( is_train=False, Y_cols=cfg.data.Y_cols, float_type=cfg.data.float_type, define_src_pos_wrt_lens=cfg.data.define_src_pos_wrt_lens, rescale_pixels=cfg.data.rescale_pixels, rescale_pixels_type=cfg.data.rescale_pixels_type, log_pixels=cfg.data.log_pixels, add_pixel_noise=cfg.data.add_pixel_noise, eff_exposure_time=cfg.data.eff_exposure_time, train_Y_mean=train_data.train_Y_mean, train_Y_std=train_data.train_Y_std, train_baobab_cfg_path=cfg.data.train_baobab_cfg_path, val_baobab_cfg_path=test_cfg.data.test_baobab_cfg_path, for_cosmology=True) master_truth = test_data.Y_df master_truth = metadata_utils.add_qphi_columns(master_truth) master_truth = metadata_utils.add_gamma_psi_ext_columns(master_truth) # Figure out how many lenses BNN will predict on (must be consecutive) if test_cfg.data.lens_indices is None: if args.lens_indices_path is None: # Test on all n_test lenses in the test set n_test = test_cfg.data.n_test lens_range = range(n_test) else: # Test on the lens indices in a text file at the specified path lens_range = [] with open(args.lens_indices_path, "r") as f: for line in f: lens_range.append(int(line.strip())) n_test = len(lens_range) print("Performing H0 inference on {:d} specified lenses...".format( n_test)) else: if args.lens_indices_path is None: # Test on the lens indices specified in the test config file lens_range = test_cfg.data.lens_indices n_test = len(lens_range) print("Performing H0 inference on {:d} specified lenses...".format( n_test)) else: raise ValueError( "Specific lens indices were specified in both the test config file and the command-line argument." ) batch_size = max(lens_range) + 1 test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, drop_last=True) # Output directory into which the H0 histograms and H0 samples will be saved out_dir = test_cfg.out_dir if not os.path.exists(out_dir): os.makedirs(out_dir) print("Destination folder path: {:s}".format(out_dir)) else: raise OSError("Destination folder already exists.") ##################### # Parameter penalty # ##################### # Instantiate original loss function with all BNN-predicted params orig_Y_cols = cfg.data.Y_cols loss_fn = getattr(h0rton.losses, cfg.model.likelihood_class)(Y_dim=test_data.Y_dim, device=device) # Not all predicted params will be sampled via MCMC params_to_remove = [] #'lens_light_R_sersic', 'src_light_R_sersic'] mcmc_Y_cols = [col for col in orig_Y_cols if col not in params_to_remove] mcmc_Y_dim = len(mcmc_Y_cols) # Instantiate loss function with just the MCMC params mcmc_loss_fn = getattr(h0rton.losses, cfg.model.likelihood_class)( Y_dim=test_data.Y_dim - len(params_to_remove), device=device) remove_param_idx, remove_idx = mcmc_utils.get_idx_for_params( mcmc_loss_fn.out_dim, orig_Y_cols, params_to_remove, cfg.model.likelihood_class) mcmc_train_Y_mean = np.delete(train_data.train_Y_mean, remove_param_idx) mcmc_train_Y_std = np.delete(train_data.train_Y_std, remove_param_idx) parameter_penalty = mcmc_utils.HybridBNNPenalty( mcmc_Y_cols, cfg.model.likelihood_class, mcmc_train_Y_mean, mcmc_train_Y_std, test_cfg.h0_posterior.exclude_velocity_dispersion, device) custom_logL_addition = parameter_penalty.evaluate null_spread = False ################### # BNN predictions # ################### # Instantiate BNN model net = getattr(h0rton.models, cfg.model.architecture)(num_classes=loss_fn.out_dim, dropout_rate=cfg.model.dropout_rate) net.to(device) # Load trained weights from saved state net, epoch = train_utils.load_state_dict_test(test_cfg.state_dict_path, net, cfg.optim.n_epochs, device) # When only generating BNN predictions (and not running MCMC), we can afford more n_dropout # otherwise, we fix n_dropout = mcmc_Y_dim + 1 if test_cfg.export.pred: n_dropout = 20 n_samples_per_dropout = test_cfg.numerics.mcmc.walkerRatio else: n_walkers = test_cfg.numerics.mcmc.walkerRatio * ( mcmc_Y_dim + 1) # (BNN params + D_dt) times walker ratio n_dropout = n_walkers // test_cfg.numerics.mcmc.walkerRatio n_samples_per_dropout = test_cfg.numerics.mcmc.walkerRatio # Initialize arrays that will store samples and BNN predictions init_pos = np.empty( [batch_size, n_dropout, n_samples_per_dropout, mcmc_Y_dim]) mcmc_pred = np.empty([batch_size, n_dropout, mcmc_loss_fn.out_dim]) with torch.no_grad(): net.train() # Send some empty forward passes through the test data without backprop to adjust batchnorm weights # (This is often not necessary. Beware if using for just 1 lens.) for nograd_pass in range(5): for X_, Y_ in test_loader: X = X_.to(device) _ = net(X) # Obtain MC dropout samples for d in range(n_dropout): net.eval() for X_, Y_ in test_loader: X = X_.to(device) Y = Y_.to(device) pred = net(X) break mcmc_pred_d = pred.cpu().numpy() # Replace BNN posterior's primary gaussian mean with truth values if test_cfg.lens_posterior_type == 'default_with_truth_mean': mcmc_pred_d[:, :len(mcmc_Y_cols)] = Y[:, :len(mcmc_Y_cols )].cpu().numpy() # Leave only the MCMC parameters in pred mcmc_pred_d = mcmc_utils.remove_parameters_from_pred( mcmc_pred_d, remove_idx, return_as_tensor=False) # Populate pred that will define the MCMC penalty function mcmc_pred[:, d, :] = mcmc_pred_d # Instantiate posterior to generate BNN samples, which will serve as initial positions for walkers bnn_post = getattr(h0rton.h0_inference.gaussian_bnn_posterior_cpu, loss_fn.posterior_name + 'CPU')( mcmc_Y_dim, mcmc_train_Y_mean, mcmc_train_Y_std) bnn_post.set_sliced_pred(mcmc_pred_d) init_pos[:, d, :, :] = bnn_post.sample( n_samples_per_dropout, sample_seed=test_cfg.global_seed + d) # contains just the lens model params, no D_dt gc.collect() # Terminate right after generating BNN predictions (no MCMC) if test_cfg.export.pred: import sys samples_path = os.path.join(out_dir, 'samples.npy') np.save(samples_path, init_pos) sys.exit() ############# # MCMC loop # ############# # Convolve MC dropout iterates with aleatoric samples init_pos = init_pos.transpose(0, 3, 1, 2).reshape( [batch_size, mcmc_Y_dim, -1]).transpose(0, 2, 1) # [batch_size, n_samples, mcmc_Y_dim] init_D_dt = np.random.uniform(0.0, 15000.0, size=(batch_size, n_walkers, 1)) pred_mean = np.mean(init_pos, axis=1) # [batch_size, mcmc_Y_dim] # Define assumed model profiles kwargs_model = dict(lens_model_list=['PEMD', 'SHEAR'], point_source_model_list=['SOURCE_POSITION'], source_light_model_list=['SERSIC_ELLIPSE']) astro_sig = test_cfg.image_position_likelihood.sigma # astrometric uncertainty # Get H0 samples for each system if not test_cfg.time_delay_likelihood.baobab_time_delays: if 'abcd_ordering_i' not in master_truth: raise ValueError( "If the time delay measurements were not generated using Baobab, the user must specify the order of image positions in which the time delays are listed, in order of increasing dec." ) kwargs_lens_eqn_solver = { 'min_distance': 0.05, 'search_window': baobab_cfg.instrument['pixel_scale'] * baobab_cfg.image['num_pix'], 'num_iter_max': 200 } total_progress = tqdm(total=n_test) realized_time_delays = pd.read_csv( test_cfg.error_model.realized_time_delays, index_col=None) # For each lens system... for i, lens_i in enumerate(lens_range): # Each lens gets a unique random state for time delay measurement error realizations. #rs_lens = np.random.RandomState(lens_i) # replaced with externally rendered time delays ########################### # Relevant data and prior # ########################### data_i = master_truth.iloc[lens_i].copy() # Set BNN pred defining parameter penalty for this lens, batch processes across n_dropout parameter_penalty.set_bnn_post_params(mcmc_pred[lens_i, :, :]) # Initialize lens model params walkers at the predictive mean init_info = dict( zip(mcmc_Y_cols, pred_mean[lens_i, :] * mcmc_train_Y_std + mcmc_train_Y_mean)) lcdm = LCDM(z_lens=data_i['z_lens'], z_source=data_i['z_src'], flat=True) true_img_dec = literal_eval(data_i['y_image']) n_img = len(true_img_dec) measured_td_sig = test_cfg.time_delay_likelihood.sigma measured_td_wrt0 = np.array( literal_eval( realized_time_delays.iloc[lens_i]['measured_td_wrt0'])) kwargs_data_joint = dict( time_delays_measured=measured_td_wrt0, time_delays_uncertainties=measured_td_sig, ) ############################# # Parameter init and bounds # ############################# lens_kwargs = mcmc_utils.get_lens_kwargs(init_info, null_spread=null_spread) ps_kwargs = mcmc_utils.get_ps_kwargs_src_plane(init_info, astro_sig) src_light_kwargs = mcmc_utils.get_light_kwargs( init_info['src_light_R_sersic'], null_spread=null_spread) special_kwargs = mcmc_utils.get_special_kwargs( n_img, astro_sig ) # image position offset and time delay distance, aka the "special" parameters kwargs_params = { 'lens_model': lens_kwargs, 'point_source_model': ps_kwargs, 'source_model': src_light_kwargs, 'special': special_kwargs, } if test_cfg.numerics.solver_type == 'NONE': solver_type = 'NONE' else: solver_type = 'PROFILE_SHEAR' if n_img == 4 else 'CENTER' #solver_type = 'NONE' kwargs_constraints = { 'num_point_source_list': [n_img], 'Ddt_sampling': True, 'solver_type': solver_type, } kwargs_likelihood = { 'time_delay_likelihood': True, 'sort_images_by_dec': True, 'prior_lens': [], 'prior_special': [], 'check_bounds': True, 'check_matched_source_position': False, 'source_position_tolerance': 0.01, 'source_position_sigma': 0.01, 'source_position_likelihood': False, 'custom_logL_addition': custom_logL_addition, 'kwargs_lens_eqn_solver': kwargs_lens_eqn_solver } ########################### # MCMC posterior sampling # ########################### fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, verbose=False, mpi=False) if i == 0: param_class = fitting_seq._updateManager.param_class n_params, param_class_Y_cols = param_class.num_param() init_pos = mcmc_utils.reorder_to_param_class( mcmc_Y_cols, param_class_Y_cols, init_pos, init_D_dt) # MCMC sample from the post-processed BNN posterior jointly with cosmology lens_i_start_time = time.time() if test_cfg.lens_posterior_type == 'default': test_cfg.numerics.mcmc.update(init_samples=init_pos[lens_i, :, :]) fitting_kwargs_list_mcmc = [['MCMC', test_cfg.numerics.mcmc]] #try: with script_utils.HiddenPrints(): chain_list_mcmc = fitting_seq.fit_sequence( fitting_kwargs_list_mcmc) kwargs_result_mcmc = fitting_seq.best_fit() lens_i_end_time = time.time() inference_time = (lens_i_end_time - lens_i_start_time) / 60.0 # min ############################# # Plotting the MCMC samples # ############################# # sampler_type : 'EMCEE' # samples_mcmc : np.array of shape `[n_mcmc_eval, n_params]` # param_mcmc : list of str of length n_params, the parameter names sampler_type, samples_mcmc, param_mcmc, _ = chain_list_mcmc[0] new_samples_mcmc = mcmc_utils.postprocess_mcmc_chain( kwargs_result_mcmc, samples_mcmc, kwargs_model, lens_kwargs[2], ps_kwargs[2], src_light_kwargs[2], special_kwargs[2], kwargs_constraints) # Plot D_dt histogram D_dt_samples = new_samples_mcmc['D_dt'].values true_D_dt = lcdm.D_dt(H_0=data_i['H0'], Om0=0.3) data_i['D_dt'] = true_D_dt # Export D_dt samples for this lens lens_inference_dict = dict( D_dt_samples=D_dt_samples, # kappa_ext=0 for these samples inference_time=inference_time, true_D_dt=true_D_dt, ) lens_inference_dict_save_path = os.path.join( out_dir, 'D_dt_dict_{0:04d}.npy'.format(lens_i)) np.save(lens_inference_dict_save_path, lens_inference_dict) # Optionally export the MCMC samples if test_cfg.export.mcmc_samples: mcmc_samples_path = os.path.join( out_dir, 'mcmc_samples_{0:04d}.csv'.format(lens_i)) new_samples_mcmc.to_csv(mcmc_samples_path, index=None) # Optionally export the D_dt histogram if test_cfg.export.D_dt_histogram: cleaned_D_dt_samples = h0_utils.remove_outliers_from_lognormal( D_dt_samples, 3) _ = plotting_utils.plot_D_dt_histogram(cleaned_D_dt_samples, lens_i, true_D_dt, save_dir=out_dir) # Optionally export the plot of MCMC chain if test_cfg.export.mcmc_chain: mcmc_chain_path = os.path.join( out_dir, 'mcmc_chain_{0:04d}.png'.format(lens_i)) plotting_utils.plot_mcmc_chain(chain_list_mcmc, mcmc_chain_path) # Optionally export posterior cornerplot of select lens model parameters with D_dt if test_cfg.export.mcmc_corner: mcmc_corner_path = os.path.join( out_dir, 'mcmc_corner_{0:04d}.png'.format(lens_i)) plotting_utils.plot_mcmc_corner( new_samples_mcmc[test_cfg.export.mcmc_cols], data_i[test_cfg.export.mcmc_cols], test_cfg.export.mcmc_col_labels, mcmc_corner_path) total_progress.update(1) gc.collect() realized_time_delays.to_csv(os.path.join(out_dir, 'realized_time_delays.csv'), index=None) total_progress.close()
fitting_seqI = FittingSequence(multi_band_listI, kwargs_modelI, kwargs_constraints, kwargs_likelihood, kwargs_params) fitting_kwargs_list = [[ 'MCMC', { 'n_burn': 100, 'n_run': 100, 'walkerRatio': 10, 'sigma_scale': .1 } ]] chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seqI.fit_sequence( fitting_kwargs_list) lens_result, source_result, lens_light_resultI, ps_result, cosmo_result = fitting_seqI.best_fit( ) fitting_seqR = FittingSequence(multi_band_listR, kwargs_modelR, kwargs_constraints, kwargs_likelihood, kwargs_params) fitting_kwargs_list = [[ 'PSO', { 'sigma_scale': 1., 'n_particles': 180, 'n_iterations': 180 } ]] chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seqR.fit_sequence( fitting_kwargs_list)
def sl_sys_analysis(): # Get command line arguments args = {} if comm_rank == 0: print(":Registered %d processes" % comm_size) args["infile"] = sys.argv[1] args["nimgs"] = sys.argv[2] args["los"] = sys.argv[3] args = comm.bcast(args) # Organize devision of strong lensing systems with open(args["infile"], "r") as myfile: limg_data = myfile.read() systems = json.loads(limg_data) sys_nr_per_proc = int(len(systems) / comm_size) start_sys = sys_nr_per_proc * comm_rank end_sys = sys_nr_per_proc * (comm_rank + 1) with open("../lens_catalogs_sie_only.json", "r") as myfile: limg_data = myfile.read() systems_prior = json.loads(limg_data) if comm_rank == 0: print("Each process will have %d systems" % sys_nr_per_proc) print("That should take app. %f min." % (sys_nr_per_proc * 20)) results = {"D_dt": []} for ii in range(len(systems))[(start_sys + 2) : end_sys]: system = systems[ii] system_prior = systems_prior[ii] print("Analysing system ID: %d" % system["losID"]) # the data set is z_lens = system_prior["zl"] z_source = 2.0 # image positions units of arcsec # time delays units of days # image brightness # amplitude (in arbitrary linear units, not magnitudes) ximg = [] yimg = [] delay = [] image_amps = [] for jj in range(system["nimgs"]): ximg.append(system["ximg"][jj]) yimg.append(system["yimg"][jj]) delay.append(system["delay"][jj]) image_amps.append(system["mags"][jj]) ximg = np.asarray(ximg) yimg = np.asarray(yimg) # 1-sigma astrometric uncertainties of the image positions # (assuming equal precision) astrometry_sigma = 0.004 delay = np.asarray(delay) d_dt = delay[1:] - delay[0] # 1-sigma uncertainties in the time-delay measurement (in units of days) d_dt_sigma = np.ones(system["nimgs"] - 1) * 2 image_amps = np.asarray(image_amps) image_amps_sigma = np.ones(system["nimgs"]) * 0.3 flux_ratios = image_amps[1:] - image_amps[0] flux_ratio_errors = np.ones(system["nimgs"] - 1) * 0.1 # lens model choicers lens_model_list = ["SPEP", "SHEAR"] fixed_lens = [] kwargs_lens_init = [] kwargs_lens_sigma = [] kwargs_lower_lens = [] kwargs_upper_lens = [] fixed_lens.append({}) kwargs_lens_init.append( { "theta_E": 1.0, "gamma": 2, "center_x": 0, "center_y": 0, "e1": 0, "e2": 0.0, } ) kwargs_lens_sigma.append( { "theta_E": 0.2, "e1": 0.1, "e2": 0.1, "gamma": 0.1, "center_x": 0.1, "center_y": 0.1, } ) kwargs_lower_lens.append( { "theta_E": 0.01, "e1": -0.5, "e2": -0.5, "gamma": 1.5, "center_x": -10, "center_y": -10, } ) kwargs_upper_lens.append( { "theta_E": 10, "e1": 0.5, "e2": 0.5, "gamma": 2.5, "center_x": 10, "center_y": 10, } ) fixed_lens.append({"ra_0": 0, "dec_0": 0}) kwargs_lens_init.append({"e1": 0.0, "e2": 0.0}) kwargs_lens_sigma.append({"e1": 0.1, "e2": 0.1}) kwargs_lower_lens.append({"e1": -0.2, "e2": -0.2}) kwargs_upper_lens.append({"e1": 0.2, "e2": 0.2}) lens_params = [ kwargs_lens_init, kwargs_lens_sigma, fixed_lens, kwargs_lower_lens, kwargs_upper_lens, ] point_source_list = ["LENSED_POSITION"] fixed_ps = [ {"ra_image": ximg, "dec_image": yimg} ] # we fix the image position coordinates kwargs_ps_init = fixed_ps kwargs_ps_sigma = [ { "ra_image": 0.01 * np.ones(len(ximg)), "dec_image": 0.01 * np.ones(len(ximg)), } ] kwargs_lower_ps = [ { "ra_image": -10 * np.ones(len(ximg)), "dec_image": -10 * np.ones(len(ximg)), } ] kwargs_upper_ps = [ {"ra_image": 10 * np.ones(len(ximg)), "dec_image": 10 * np.ones(len(ximg))} ] ps_params = [ kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps, ] # we let some freedome in how well the actual image positions are matching those # given by the data (indicated as 'ra_image', 'dec_image' and held fixed while fitting) fixed_cosmo = {} kwargs_cosmo_init = { "D_dt": 5000, "delta_x_image": np.zeros_like(ximg), "delta_y_image": np.zeros_like(ximg), } kwargs_cosmo_sigma = { "D_dt": 10000, "delta_x_image": np.ones_like(ximg) * astrometry_sigma, "delta_y_image": np.ones_like(ximg) * astrometry_sigma, } kwargs_lower_cosmo = { "D_dt": 0, "delta_x_image": np.ones_like(ximg) * (-1), "delta_y_image": np.ones_like(ximg) * (-1), } kwargs_upper_cosmo = { "D_dt": 10000, "delta_x_image": np.ones_like(ximg) * (1), "delta_y_image": np.ones_like(ximg) * (1), } cosmo_params = [ kwargs_cosmo_init, kwargs_cosmo_sigma, fixed_cosmo, kwargs_lower_cosmo, kwargs_upper_cosmo, ] kwargs_params = { "lens_model": lens_params, "point_source_model": ps_params, "cosmography": cosmo_params, } # ## setup options for likelihood and parameter sampling kwargs_constraints = { "num_point_source_list": [4], # any proposed lens model must satisfy the image positions # appearing at the position of the point sources being sampeld "solver_type": "PROFILE_SHEAR", "cosmo_type": "D_dt", # sampling of the time-delay distance # explicit modelling of the astrometric imperfection of # the point source positions "point_source_offset": True, } kwargs_likelihood = { "check_bounds": True, "point_source_likelihood": True, "position_uncertainty": astrometry_sigma, "check_solver": True, "solver_tolerance": 0.001, "time_delay_likelihood": True, "image_likelihood": False, # this needs to be explicitly given when not having imaging data "flux_ratio_likelihood": True, # enables the flux ratio likelihood } kwargs_data_joint = { "time_delays_measured": d_dt, "time_delays_uncertainties": d_dt_sigma, "flux_ratios": flux_ratios, "flux_ratio_errors": flux_ratio_errors, } kwargs_model = { "lens_model_list": lens_model_list, "point_source_model_list": point_source_list, } mpi = False # MPI possible, but not supported through that notebook. fitting_seq = FittingSequence( kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, ) fitting_kwargs_list = [ # ['update_settings', {'lens_add_fixed': [[0, ['gamma']]]}], ["PSO", {"sigma_scale": 1.0, "n_particles": 100, "n_iterations": 100}] ] 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_result = ( fitting_seq.best_fit() ) # and now we run the MCMC fitting_kwargs_list = [ ["PSO", {"sigma_scale": 0.1, "n_particles": 100, "n_iterations": 100}], [ "MCMC", {"n_burn": 200, "n_run": 200, "walkerRatio": 10, "sigma_scale": 0.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_result = ( fitting_seq.best_fit() ) # print("number of non-linear parameters in the MCMC process: ", len(param_mcmc)) # print("parameters in order: ", param_mcmc) print("number of evaluations in the MCMC process: ", np.shape(samples_mcmc)[0]) param = Param( kwargs_model, fixed_lens, kwargs_fixed_ps=fixed_ps, kwargs_fixed_cosmo=fixed_cosmo, kwargs_lens_init=lens_result, **kwargs_constraints, ) # the number of non-linear parameters and their names # num_param, param_list = param.num_param() lensAnalysis = LensAnalysis(kwargs_model) mcmc_new_list = [] labels_new = [ r"$\gamma$", r"$\phi_{ext}$", r"$\gamma_{ext}$", r"$D_{\Delta t}$", ] for i in range(len(samples_mcmc)): # transform the parameter position of the MCMC chain in a # lenstronomy convention with keyword arguments # kwargs_lens_out, kwargs_light_source_out, kwargs_light_lens_out, kwargs_ps_out, kwargs_cosmo = param.args2kwargs( samples_mcmc[i] ) D_dt = kwargs_cosmo["D_dt"] gamma = kwargs_lens_out[0]["gamma"] phi_ext, gamma_ext = lensAnalysis._lensModelExtensions.external_shear( kwargs_lens_out ) mcmc_new_list.append([gamma, phi_ext, gamma_ext, D_dt]) # plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True) # and here the predicted angular diameter distance from a # default cosmology (attention for experimenter bias!) lensCosmo = LensCosmo(z_lens=z_lens, z_source=z_source) results["D_dt"].append(D_dt) f = h5py.File( "H0_" + args["los"] + "_nimgs" + str(args["nimgs"]) + "_" + str(comm_rank) + ".hdf5", "w", ) hf.create_dataset("D_dt", data=results["D_dt"]) hf.close()