def postprocess_mcmc_chain(kwargs_result, samples, kwargs_model, fixed_lens_kwargs, fixed_ps_kwargs, fixed_src_light_kwargs, fixed_special_kwargs, kwargs_constraints, kwargs_fixed_lens_light=None, verbose=False, forward_modeling=False): """Postprocess the MCMC chain for making the chains consistent with the optimized lens model and converting parameters Returns ------- pandas.DataFrame processed MCMC chain, where each row is a sample """ param = Param(kwargs_model, fixed_lens_kwargs, kwargs_fixed_ps=fixed_ps_kwargs, kwargs_fixed_source=fixed_src_light_kwargs, kwargs_fixed_special=fixed_special_kwargs, kwargs_fixed_lens_light=kwargs_fixed_lens_light, kwargs_lens_init=kwargs_result['kwargs_lens'], **kwargs_constraints) if verbose: param.print_setting() n_samples = len(samples) processed = [] for i in range(n_samples): kwargs = {} kwargs_out = param.args2kwargs(samples[i]) kwargs_lens_out, kwargs_special_out, kwargs_ps_out, kwargs_source_out, kwargs_lens_light_out = kwargs_out['kwargs_lens'], kwargs_out['kwargs_special'], kwargs_out['kwargs_ps'], kwargs_out['kwargs_source'], kwargs_out['kwargs_lens_light'] for k, v in kwargs_lens_out[0].items(): kwargs['lens_mass_{:s}'.format(k)] = v for k, v in kwargs_lens_out[1].items(): kwargs['external_shear_{:s}'.format(k)] = v if forward_modeling: for k, v in kwargs_source_out[0].items(): kwargs['src_light_{:s}'.format(k)] = v for k, v in kwargs_lens_light_out[0].items(): kwargs['lens_light_{:s}'.format(k)] = v else: kwargs['src_light_R_sersic'] = kwargs_source_out[0]['R_sersic'] if 'ra_source' in kwargs_ps_out[0]: kwargs['src_light_center_x'] = kwargs_ps_out[0]['ra_source'] - kwargs_lens_out[0]['center_x'] kwargs['src_light_center_y'] = kwargs_ps_out[0]['dec_source'] - kwargs_lens_out[0]['center_y'] for k, v in kwargs_special_out.items(): kwargs[k] = v processed.append(kwargs) processed_df = pd.DataFrame(processed) processed_df = metadata_utils.add_qphi_columns(processed_df) processed_df = metadata_utils.add_gamma_psi_ext_columns(processed_df) return processed_df
class TestParam(object): def setup(self): kwargs_model = { 'lens_model_list': ['SPEP'], 'source_light_model_list': ['GAUSSIAN'], 'lens_light_model_list': ['SERSIC'], 'point_source_model_list': ['LENSED_POSITION'], 'multi_plane': True, 'lens_redshift_list': [0.5], 'z_source': 2, 'source_redshift_list': [0.5] } kwargs_param = { 'num_point_source_list': [2], 'lens_redshift_sampling_indexes': [0], 'source_redshift_sampling_indexes': [0], 'image_plane_source_list': [True] } kwargs_fixed_lens = [{'gamma': 1.9}] # for SPEP lens kwargs_fixed_source = [{ 'sigma': 0.1, 'center_x': 0.2, 'center_y': 0.2 }] kwargs_fixed_ps = [{'ra_image': [-1, 1], 'dec_image': [-1, 1]}] kwargs_fixed_lens_light = [{}] kwargs_fixed_cosmo = [{}] self.param_class = Param( kwargs_model, kwargs_fixed_lens=kwargs_fixed_lens, kwargs_fixed_source=kwargs_fixed_source, kwargs_fixed_lens_light=kwargs_fixed_lens_light, kwargs_fixed_ps=kwargs_fixed_ps, kwargs_fixed_special=kwargs_fixed_cosmo, **kwargs_param) self.param_class.print_setting() def test_num_param(self): num_param, list = self.param_class.num_param() assert list[0] == 'theta_E_lens0' assert num_param == 10 kwargs_model = { 'lens_model_list': ['SPEP'], 'source_light_model_list': ['GAUSSIAN'], 'lens_light_model_list': ['SERSIC'], 'point_source_model_list': ['LENSED_POSITION'] } kwargs_param = {} kwargs_fixed_lens = [{'gamma': 1.9}] # for SPEP lens kwargs_fixed_source = [{ 'sigma': 0.1, 'center_x': 0.2, 'center_y': 0.2 }] kwargs_fixed_ps = [{'ra_image': [-1, 1], 'dec_image': [-1, 1]}] kwargs_fixed_lens_light = [{}] kwargs_fixed_cosmo = [{}] param_class_linear = Param(kwargs_model, kwargs_fixed_lens, kwargs_fixed_source, kwargs_fixed_lens_light, kwargs_fixed_ps, kwargs_fixed_cosmo, linear_solver=True, **kwargs_param) num_param, list = param_class_linear.num_param() assert list[0] == 'theta_E_lens0' print(list) assert len(list) == num_param assert num_param == 9 def test_num_param_linear(self): num_param = self.param_class.num_param_linear() assert num_param == 4 def test_get_params(self): kwargs_true_lens = [{ 'theta_E': 1., 'gamma': 1.9, 'e1': 0.01, 'e2': -0.01, 'center_x': 0., 'center_y': 0. }] #for SPEP lens kwargs_true_source = [{ 'amp': 1 * 2 * np.pi * 0.1**2, 'center_x': 0.2, 'center_y': 0.2, 'sigma': 0.1 }] kwargs_true_lens_light = [{ 'center_x': -0.06, 'center_y': 0.4, 'phi_G': 4.8, 'q': 0.86, 'n_sersic': 1.7, 'amp': 11.8, 'R_sersic': 0.697, 'phi_G_2': 0 }] kwargs_true_ps = [{ 'point_amp': [1, 1], 'ra_image': [-1, 1], 'dec_image': [-1, 1] }] kwargs_special = {'z_sampling': [0.5]} args = self.param_class.kwargs2args( kwargs_true_lens, kwargs_true_source, kwargs_lens_light=kwargs_true_lens_light, kwargs_ps=kwargs_true_ps, kwargs_special=kwargs_special) kwargs_return = self.param_class.args2kwargs(args) lens_dict_list = kwargs_return['kwargs_lens'] lens_light_dict_list = kwargs_return['kwargs_lens_light'] lens_dict = lens_dict_list[0] assert lens_dict['theta_E'] == 1. assert lens_dict['gamma'] == 1.9 assert lens_dict['e1'] == 0.01 assert lens_dict['e2'] == -0.01 assert lens_dict['center_x'] == 0. assert lens_dict['center_y'] == 0. assert lens_light_dict_list[0]['center_x'] == -0.06 def test_get_cosmo(self): kwargs_model = { 'lens_model_list': ['SPEP'], 'source_light_model_list': ['GAUSSIAN'], 'lens_light_model_list': ['SERSIC'], 'point_source_model_list': ['LENSED_POSITION'], } kwargs_param = {'Ddt_sampling': True} kwargs_fixed_lens = [{'gamma': 1.9}] # for SPEP lens kwargs_fixed_source = [{ 'sigma': 0.1, 'center_x': 0.2, 'center_y': 0.2 }] kwargs_fixed_ps = [{'ra_image': [-1, 1], 'dec_image': [-1, 1]}] kwargs_fixed_lens_light = [{}] kwargs_fixed_cosmo = {'D_dt': 1000} param_class = Param(kwargs_model, kwargs_fixed_lens=kwargs_fixed_lens, kwargs_fixed_source=kwargs_fixed_source, kwargs_fixed_lens_light=kwargs_fixed_lens_light, kwargs_fixed_ps=kwargs_fixed_ps, kwargs_fixed_special=kwargs_fixed_cosmo, **kwargs_param) kwargs_true_lens = [{ 'theta_E': 1., 'gamma': 1.9, 'e1': 0.01, 'e2': -0.01, 'center_x': 0., 'center_y': 0. }] # for SPEP lens kwargs_true_source = [{ 'amp': 1 * 2 * np.pi * 0.1**2, 'center_x': 0.2, 'center_y': 0.2, 'sigma': 0.1 }] kwargs_true_lens_light = [{ 'center_x': -0.06, 'center_y': 0.4, 'phi_G': 4.8, 'q': 0.86, 'n_sersic': 1.7, 'amp': 11.8, 'R_sersic': 0.697, 'phi_G_2': 0 }] kwargs_true_ps = [{ 'point_amp': [1, 1], 'ra_image': [-1, 1], 'dec_image': [-1, 1] }] args = param_class.kwargs2args( kwargs_true_lens, kwargs_true_source, kwargs_lens_light=kwargs_true_lens_light, kwargs_ps=kwargs_true_ps, kwargs_special={'D_dt': 1000}) assert param_class.specialParams._D_dt_sampling is True def test_mass_scaling(self): kwargs_model = { 'lens_model_list': ['SIS', 'NFW', 'NFW', 'SIS', 'SERSIC', 'HERNQUIST'] } kwargs_constraints = {'mass_scaling_list': [False, 1, 1, 1, 1, 1]} kwargs_fixed_lens = [{}, { 'alpha_Rs': 0.1 }, { 'alpha_Rs': 0.3 }, { 'theta_E': 0.1 }, { 'k_eff': 0.3 }, { 'sigma0': 1 }] kwargs_fixed_cosmo = {} param_class = Param(kwargs_model, kwargs_fixed_lens=kwargs_fixed_lens, kwargs_fixed_special=kwargs_fixed_cosmo, **kwargs_constraints) kwargs_lens = [{ 'theta_E': 1, 'center_x': 0, 'center_y': 0 }, { 'alpha_Rs': 0.1, 'Rs': 5, 'center_x': 1., 'center_y': 0 }, { 'alpha_Rs': 0.1, 'Rs': 5, 'center_x': 0, 'center_y': 1. }, { 'theta_E': 0.1, 'center_x': 3, 'center_y': 1. }, { 'k_eff': 0.3, 'R_sersic': 1, 'n_sersic': 2, 'center_x': 3, 'center_y': 1. }, { 'sigma0': 1, 'Rs': 1, 'center_x': 3, 'center_y': 1. }] kwargs_source = [] kwargs_lens_light = [] kwargs_ps = [] mass_scale = 2 kwargs_cosmo = {'scale_factor': [mass_scale]} args = param_class.kwargs2args(kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps, kwargs_special=kwargs_cosmo) assert args[-1] == mass_scale kwargs_return = param_class.args2kwargs(args) kwargs_lens = kwargs_return['kwargs_lens'] print(kwargs_lens, 'test') assert kwargs_lens[0]['theta_E'] == 1 assert kwargs_lens[1]['alpha_Rs'] == 0.1 * mass_scale assert kwargs_lens[2]['alpha_Rs'] == 0.3 * mass_scale assert kwargs_lens[3]['theta_E'] == 0.1 * mass_scale assert kwargs_lens[4]['k_eff'] == 0.3 * mass_scale assert kwargs_lens[5]['sigma0'] == 1 * mass_scale kwargs_return = param_class.args2kwargs(args, bijective=True) kwargs_lens = kwargs_return['kwargs_lens'] assert kwargs_lens[0]['theta_E'] == 1 assert kwargs_lens[1]['alpha_Rs'] == 0.1 assert kwargs_lens[2]['alpha_Rs'] == 0.3 def test_joint_lens_with_light(self): kwargs_model = { 'lens_model_list': ['CHAMELEON'], 'lens_light_model_list': ['CHAMELEON'] } i_light, k_lens = 0, 0 kwargs_constraints = { 'joint_lens_with_light': [[ i_light, k_lens, ['w_t', 'w_c', 'center_x', 'center_y', 'e1', 'e2'] ]] } kwargs_lens = [{'alpha_1': 10}] kwargs_lens_light = [{ 'amp': 1, 'w_t': 0.5, 'w_c': 0.1, 'center_x': 0, 'center_y': 0.3, 'e1': 0.1, 'e2': -0.2 }] param = Param(kwargs_model=kwargs_model, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_lens_light=kwargs_lens_light) kwargs_return = param.args2kwargs(args) kwargs_lens_out = kwargs_return['kwargs_lens'] kwargs_lens_light_out = kwargs_return['kwargs_lens_light'] assert kwargs_lens_out[0]['w_c'] == kwargs_lens_light[0]['w_c'] assert kwargs_lens_light_out[0]['w_c'] == kwargs_lens_light[0]['w_c'] kwargs_model = { 'lens_model_list': ['SIS'], 'lens_light_model_list': ['SERSIC'] } i_light, k_lens = 0, 0 kwargs_constraints = { 'joint_lens_with_light': [[i_light, k_lens, ['center_x', 'center_y']]] } # list[[i_point_source, k_source, ['param_name1', 'param_name2', ...]], [ kwargs_lens = [{'theta_E': 1, 'center_x': 0, 'center_y': 0}] kwargs_lens_light = [{ 'amp': 1, 'R_sersic': 0.5, 'n_sersic': 2, 'center_x': 1, 'center_y': 1 }] param = Param(kwargs_model=kwargs_model, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_lens_light=kwargs_lens_light) #kwargs_lens_out, kwargs_source_out, _, kwargs_ps_out, _ = param.args2kwargs(args) kwargs_return = param.args2kwargs(args) kwargs_lens_out = kwargs_return['kwargs_lens'] assert kwargs_lens_out[0]['theta_E'] == kwargs_lens[0]['theta_E'] assert kwargs_lens_out[0]['center_x'] == kwargs_lens_light[0][ 'center_x'] def test_joint_lens_with_light_dict(self): kwargs_model = { 'lens_model_list': ['SHEAR'], 'lens_light_model_list': ['SERSIC'] } i_light, k_lens = 0, 0 kwargs_constraints = { 'joint_lens_with_light': [[i_light, k_lens, { 'ra_0': 'center_x', 'dec_0': 'center_y' }]] } kwargs_lens = [{'gamma1': 0.05, 'gamma2': 0.06}] kwargs_lens_light = [{ 'amp': 1, 'R_sersic': 1, 'n_sersic': 4, 'center_x': 0.1, 'center_y': 0.3 }] param = Param(kwargs_model=kwargs_model, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_lens_light=kwargs_lens_light) kwargs_return = param.args2kwargs(args) kwargs_lens_out = kwargs_return['kwargs_lens'] kwargs_lens_light_out = kwargs_return['kwargs_lens_light'] assert kwargs_lens_out[0]['gamma1'] == kwargs_lens[0]['gamma1'] assert kwargs_lens_out[0]['ra_0'] == kwargs_lens_light[0]['center_x'] assert kwargs_lens_out[0]['dec_0'] == kwargs_lens_light[0]['center_y'] def test_joint_source_with_point_source(self): kwargs_model = { 'lens_model_list': ['SIS'], 'source_light_model_list': ['SERSIC'], 'point_source_model_list': ['SOURCE_POSITION'] } i_source, k_ps = 0, 0 kwargs_constraints = { 'joint_source_with_point_source': [[k_ps, i_source]] } # list[[i_point_source, k_source, ['param_name1', 'param_name2', ...]], [ kwargs_lens = [{'theta_E': 1, 'center_x': 0, 'center_y': 0}] kwargs_source = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 1, 'center_y': 1 }] kwargs_ps = [{'ra_source': 0.5, 'dec_source': 0.5}] param = Param(kwargs_model=kwargs_model, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_source=kwargs_source, kwargs_ps=kwargs_ps) kwargs_return = param.args2kwargs(args) kwargs_lens_out = kwargs_return['kwargs_lens'] kwargs_source_out = kwargs_return['kwargs_source'] #kwargs_lens_out, kwargs_source_out, _, kwargs_ps_out, _ = param.args2kwargs(args) assert kwargs_lens_out[0]['theta_E'] == kwargs_lens[0]['theta_E'] assert kwargs_source_out[0]['center_x'] == kwargs_ps[0]['ra_source'] kwargs_model = { 'lens_model_list': ['SIS'], 'source_light_model_list': ['SERSIC'], 'point_source_model_list': ['LENSED_POSITION'] } i_source, k_ps = 0, 0 kwargs_constraints = { 'joint_source_with_point_source': [[k_ps, i_source]] } # list[[i_point_source, k_source, ['param_name1', 'param_name2', ...]], [ kwargs_lens = [{'theta_E': 1, 'center_x': 0, 'center_y': 0}] kwargs_source = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 1, 'center_y': 1 }] kwargs_ps = [{'ra_image': [0.5], 'dec_image': [0.5]}] param = Param(kwargs_model=kwargs_model, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_source=kwargs_source, kwargs_ps=kwargs_ps) #kwargs_lens_out, kwargs_source_out, _, kwargs_ps_out, _ = param.args2kwargs(args) kwargs_return = param.args2kwargs(args) kwargs_lens_out = kwargs_return['kwargs_lens'] kwargs_source_out = kwargs_return['kwargs_source'] assert kwargs_lens_out[0]['theta_E'] == kwargs_lens[0]['theta_E'] npt.assert_almost_equal(kwargs_source_out[0]['center_x'], -0.207, decimal=2) def test_joint_lens_light_with_point_source(self): kwargs_model = { 'lens_model_list': ['SIS'], 'source_light_model_list': ['SERSIC'], 'point_source_model_list': ['LENSED_POSITION'], 'lens_light_model_list': ['SERSIC'] } i_lens_light, k_ps = 0, 0 kwargs_constraints = { 'joint_lens_light_with_point_source': [[k_ps, i_lens_light]] } # list[[i_point_source, k_source, ['param_name1', 'param_name2', ...]], [ kwargs_lens = [{'theta_E': 1, 'center_x': 0, 'center_y': 0}] kwargs_source = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 1, 'center_y': 1 }] kwargs_lens_light = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 0.2, 'center_y': 0.2 }] kwargs_ps = [{'ra_image': [0.5], 'dec_image': [0.5]}] param = Param(kwargs_model=kwargs_model, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_source=kwargs_source, kwargs_lens_light=kwargs_lens_light, kwargs_ps=kwargs_ps) kwargs_return = param.args2kwargs(args) kwargs_lens_light_out = kwargs_return['kwargs_lens_light'] kwargs_ps_out = kwargs_return['kwargs_ps'] assert kwargs_lens_light_out[0]['center_x'] == kwargs_ps_out[0][ 'ra_image'] def test_logsampling(self): kwargs_model = { 'lens_model_list': ['SIS'], 'source_light_model_list': ['SERSIC'], 'point_source_model_list': ['LENSED_POSITION'], 'lens_light_model_list': ['SERSIC'] } kwargs_constraints = {'log_sampling_lens': [[0, ['theta_E']]]} kwargs_lens = [{'theta_E': 0.1, 'center_x': 0, 'center_y': 0}] kwargs_source = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 1, 'center_y': 1 }] kwargs_lens_light = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 0.2, 'center_y': 0.2 }] kwargs_ps = [{'ra_image': [0.5], 'dec_image': [0.5]}] param = Param(kwargs_model=kwargs_model, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_source=kwargs_source, kwargs_lens_light=kwargs_lens_light, kwargs_ps=kwargs_ps) kwargs_return = param.args2kwargs(args) kwargs_lens_out = kwargs_return['kwargs_lens'] assert args[0] == -1 assert kwargs_lens_out[0]['theta_E'] == 0.1 def test_with_solver(self): kwargs_model = { 'lens_model_list': ['SPEP'], 'source_light_model_list': ['SERSIC'], 'point_source_model_list': ['LENSED_POSITION'] } i_lens_light, k_ps = 0, 0 kwargs_constraints = { 'solver_type': 'PROFILE', 'num_point_source_list': [4] } kwargs_lens = [{ 'theta_E': 1, 'gamma': 2, 'e1': 0.1, 'e2': 0.1, 'center_x': 0, 'center_y': 0 }] kwargs_source = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 1, 'center_y': 1 }] kwargs_lens_light = [{ 'amp': 1, 'n_sersic': 2, 'R_sersic': 0.3, 'center_x': 0.2, 'center_y': 0.2 }] lensModel = LensModel(lens_model_list=['SPEP']) lensEquationSlover = LensEquationSolver(lensModel=lensModel) x_image, y_image = lensEquationSlover.image_position_from_source( sourcePos_x=0.0, sourcePos_y=0.01, kwargs_lens=kwargs_lens) print(x_image, y_image, 'test') kwargs_ps = [{'ra_image': x_image, 'dec_image': y_image}] param = Param(kwargs_model=kwargs_model, kwargs_lens_init=kwargs_lens, **kwargs_constraints) args = param.kwargs2args(kwargs_lens=kwargs_lens, kwargs_source=kwargs_source, kwargs_lens_light=kwargs_lens_light, kwargs_ps=kwargs_ps) #kwargs_lens_out, kwargs_source_out, kwargs_lens_light_out, kwargs_ps_out, _ = param.args2kwargs(args) kwargs_return = param.args2kwargs(args) kwargs_lens_out = kwargs_return['kwargs_lens'] kwargs_ps_out = kwargs_return['kwargs_ps'] dist = param.check_solver(kwargs_lens=kwargs_lens_out, kwargs_ps=kwargs_ps_out) npt.assert_almost_equal(dist, 0, decimal=10) def test_num_point_source_images(self): num = self.param_class.num_point_source_images assert num == 2 def test_shapelet_lens(self): kwargs_model = { 'lens_model_list': ['SHAPELETS_CART'], 'source_light_model_list': [], 'lens_light_model_list': [], 'point_source_model_list': [] } kwargs_param = {'num_shapelet_lens': 6} kwargs_fixed_lens = [{'beta': 1}] # for SPEP lens kwargs_fixed_source = [{}] kwargs_fixed_ps = [{}] kwargs_fixed_lens_light = [{}] kwargs_fixed_cosmo = [{}] self.param_class = Param( kwargs_model, kwargs_fixed_lens=kwargs_fixed_lens, kwargs_fixed_source=kwargs_fixed_source, kwargs_fixed_lens_light=kwargs_fixed_lens_light, kwargs_fixed_ps=kwargs_fixed_ps, kwargs_fixed_special=kwargs_fixed_cosmo, **kwargs_param) self.param_class.print_setting()