def test_solver_simplified(self): lens_model_list = ['SPEP', 'SHEAR_GAMMA_PSI'] lensModel = LensModel(lens_model_list) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = -0.1 deltapix = 0.05 numPix = 150 gamma = 1.9 gamma_ext = 0.05 psi_ext = 0.4 #e1, e2 = param_util.phi_gamma_ellipticity(phi=psi_ext, gamma=gamma_ext) kwargs_lens = [{'theta_E': 1., 'gamma': gamma, 'e1': 0.1, 'e2': -0.1, 'center_x': 0.1, 'center_y': -0.1}, {'gamma_ext': gamma_ext, 'psi_ext': psi_ext}] x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) e1_new, e2_new = param_util.phi_gamma_ellipticity(phi=0., gamma=gamma_ext+0.1) kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0}, {'gamma_ext': gamma_ext + 0.1, 'psi_ext': 0}] solver = Solver4Point(lensModel, solver_type='PROFILE_SHEAR') kwargs_lens_new, accuracy = solver.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init) assert accuracy < 10**(-10) x_source, y_source = lensModel.ray_shooting(x_pos, y_pos, kwargs_lens_new) x_source, y_source = np.mean(x_source), np.mean(y_source) x_pos_new, y_pos_new = lensEquationSolver.findBrightImage(x_source, y_source, kwargs_lens_new, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) print(x_pos, x_pos_new) x_pos = np.sort(x_pos) x_pos_new = np.sort(x_pos_new) y_pos = np.sort(y_pos) y_pos_new = np.sort(y_pos_new) for i in range(len(x_pos)): npt.assert_almost_equal(x_pos[i], x_pos_new[i], decimal=6) npt.assert_almost_equal(y_pos[i], y_pos_new[i], decimal=6)
def test_shapelet_cart(self): lens_model_list = ['SHAPELETS_CART', 'SIS'] lens = LensModel(lens_model_list) solver = Solver2Point(lens, solver_type='SHAPELETS') image_position = LensEquationSolver(lens) sourcePos_x = 0.1 sourcePos_y = 0.03 deltapix = 0.05 numPix = 100 kwargs_lens = [{ 'coeffs': [1., 0., 0.1, 1.], 'beta': 1. }, { 'theta_E': 1., 'center_x': -0.1, 'center_y': 0.1 }] x_pos, y_pos = image_position.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=2, min_distance=deltapix, search_window=numPix * deltapix, precision_limit=10**(-10)) print(len(x_pos), 'number of images') x_pos = x_pos[:2] y_pos = y_pos[:2] kwargs_init = [{ 'coeffs': [1., 0., 0.1, 1.], 'beta': 1. }, { 'theta_E': 1., 'center_x': -0.1, 'center_y': 0.1 }] kwargs_out, precision = solver.constraint_lensmodel( x_pos, y_pos, kwargs_init) print(kwargs_out, 'output') source_x, source_y = lens.ray_shooting(x_pos[0], y_pos[0], kwargs_out) x_pos_new, y_pos_new = image_position.findBrightImage( source_x, source_y, kwargs_out, numImages=2, min_distance=deltapix, search_window=numPix * deltapix) npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=3) npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=3) npt.assert_almost_equal(kwargs_out[0]['coeffs'][1], kwargs_lens[0]['coeffs'][1], decimal=3) npt.assert_almost_equal(kwargs_out[0]['coeffs'][2], kwargs_lens[0]['coeffs'][2], decimal=3)
def test_solver_spep(self): lens_model_list = ['SPEP'] lensModel = LensModel(lens_model_list) solver = Solver4Point(lensModel) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = -0.1 deltapix = 0.05 numPix = 150 gamma = 1.9 phi_G, q = 0.5, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi_G, q) kwargs_lens = [{'theta_E': 1., 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0.1, 'center_y': -0.1}] x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix*deltapix) phi_G, q = 1.5, 0.9 e1, e2 = param_util.phi_q2_ellipticity(phi_G, q) kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0., 'center_y': 0}] kwargs_lens_new, accuracy = solver.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], kwargs_lens[0]['theta_E'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['e1'], kwargs_lens[0]['e1'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['e2'], kwargs_lens[0]['e2'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_x'], kwargs_lens[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_y'], kwargs_lens[0]['center_y'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], 1., decimal=3) lensModel = LensModel(lens_model_list=lens_model_list) x_source_new, y_source_new = lensModel.ray_shooting(x_pos, y_pos, kwargs_lens_new) dist = np.sqrt((x_source_new - x_source_new[0]) ** 2 + (y_source_new - y_source_new[0]) ** 2) print(dist) assert np.max(dist) < 0.000001
def solve_leq(self, xsrc, ysrc, lensmodel, lens_model_params, brightimg, precision_lim=10**(-10), nitermax=10): lensEquationSolver = LensEquationSolver(lensModel=lensmodel) if brightimg: x_image, y_image = lensEquationSolver.findBrightImage( xsrc, ysrc, lens_model_params, arrival_time_sort=False, min_distance=0.02, search_window=3.5, precision_limit=precision_lim, num_iter_max=nitermax) else: x_image, y_image = lensEquationSolver.image_position_from_source( kwargs_lens=lens_model_params, sourcePos_x=xsrc, sourcePos_y=ysrc, min_distance=self.min_distance, search_window=self.search_window, precision_limit=self.precision_limit, num_iter_max=self.num_iter_max, arrival_time_sort=False) return x_image, y_image
class MockLensBase(object): def __init__(self, source_x, source_y, kwargs): self.lensModel = LensModel(['SPEMD', 'SHEAR']) self.solver = LensEquationSolver(self.lensModel) self.x_image, self.y_image = self.solver.findBrightImage( source_x, source_y, kwargs)
def test_plot_quasar_images(self): lens_model_list = ['EPL', 'SHEAR'] z_source = 1.5 kwargs_lens = [{ 'theta_E': 1., 'gamma': 2., 'e1': 0.02, 'e2': -0.09, 'center_x': 0, 'center_y': 0 }, { 'gamma1': 0.01, 'gamma2': 0.03 }] lensmodel = LensModel(lens_model_list) solver = LensEquationSolver(lensmodel) source_x, source_y = 0.07, 0.03 x_image, y_image = solver.findBrightImage(source_x, source_y, kwargs_lens) source_fwhm_parsec = 40. plot_quasar_images(lensmodel, x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source) plt.close()
def test_constraint_lensmodel(self): lens_model_list = ['SPEP', 'SIS'] lensModel = LensModel(lens_model_list) solver = Solver(solver_type='PROFILE', lensModel=lensModel, num_images=4) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = -0.1 deltapix = 0.05 numPix = 150 gamma = 1.9 kwargs_lens = [{'theta_E': 1., 'gamma': gamma,'q': 0.8, 'phi_G': 0.5, 'center_x': 0.1, 'center_y': -0.1}, {'theta_E': 0.1, 'center_x': 0.5, 'center_y': 0}] x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix*deltapix) kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma,'q': 0.9, 'phi_G': 1.5, 'center_x': 0., 'center_y': 0}, {'theta_E': 0.1, 'center_x': 0.5, 'center_y': 0}] kwargs_lens_new, accuracy = solver.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], kwargs_lens[0]['theta_E'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['q'], kwargs_lens[0]['q'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['phi_G'], kwargs_lens[0]['phi_G'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_x'], kwargs_lens[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_y'], kwargs_lens[0]['center_y'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], 1., decimal=3) lensModel = LensModel(lens_model_list=lens_model_list) x_source_new, y_source_new = lensModel.ray_shooting(x_pos, y_pos, kwargs_lens_new) dist = np.sqrt((x_source_new - x_source_new[0]) ** 2 + (y_source_new - y_source_new[0]) ** 2) assert np.max(dist) < 0.000001 kwargs_ps4 = [{'ra_image': x_pos, 'dec_image': y_pos}] kwargs_lens_new = solver.update_solver(kwargs_lens_init, kwargs_ps4) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], kwargs_lens[0]['theta_E'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['q'], kwargs_lens[0]['q'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['phi_G'], kwargs_lens[0]['phi_G'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_x'], kwargs_lens[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_y'], kwargs_lens[0]['center_y'], decimal=3)
def test_solver_nfw(self): lens_model_list = ['NFW_ELLIPSE', 'SIS'] lensModel = LensModel(lens_model_list) solver = Solver4Point(lensModel) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = -0.1 deltapix = 0.05 numPix = 150 Rs = 4. phi_G, q = 0.5, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi_G, q) kwargs_lens = [{ 'theta_Rs': 1., 'Rs': Rs, 'e1': e1, 'e2': e2, 'center_x': 0.1, 'center_y': -0.1 }, { 'theta_E': 1, 'center_x': 0, 'center_y': 0 }] x_pos, y_pos = lensEquationSolver.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) phi_G, q = 1.5, 0.9 e1, e2 = param_util.phi_q2_ellipticity(phi_G, q) kwargs_lens_init = [{ 'theta_Rs': 0.5, 'Rs': Rs, 'e1': e1, 'e2': e2, 'center_x': 0., 'center_y': 0 }, kwargs_lens[1]] kwargs_lens_new, accuracy = solver.constraint_lensmodel( x_pos, y_pos, kwargs_lens_init) npt.assert_almost_equal(kwargs_lens_new[0]['theta_Rs'], kwargs_lens[0]['theta_Rs'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['e1'], kwargs_lens[0]['e1'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['e2'], kwargs_lens[0]['e2'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_x'], kwargs_lens[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_y'], kwargs_lens[0]['center_y'], decimal=3)
def test_solver_profile_shear_2(self): lens_model_list = ['SPEP', 'SHEAR'] lensModel = LensModel(lens_model_list) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0. sourcePos_y = 0.1 deltapix = 0.05 numPix = 150 gamma = 1.98 e1, e2 = -0.04, -0.01 kwargs_shear = {'e1': e1, 'e2': e2} # shear values to the source plane kwargs_spemd = {'theta_E': 1.66, 'gamma': gamma, 'center_x': 0.0, 'center_y': 0.0, 'e1': 0.1, 'e2': 0.05} # parameters of the deflector lens model kwargs_lens = [kwargs_spemd, kwargs_shear] x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) print(x_pos, y_pos, 'test positions') gamma_ext = np.sqrt(e1 ** 2 + e2 ** 2) e1_init, e2_init = param_util.phi_gamma_ellipticity(gamma=gamma_ext, phi=-1.3) kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': 0, 'e2': 0, 'center_x': 0., 'center_y': 0}, {'e1': e1_init, 'e2': e2_init}] solver = Solver4Point(lensModel, solver_type='PROFILE_SHEAR') kwargs_lens_new, accuracy = solver.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init) assert accuracy < 10**(-10) x_source, y_source = lensModel.ray_shooting(x_pos, y_pos, kwargs_lens_new) x_source, y_source = np.mean(x_source), np.mean(y_source) x_pos_new, y_pos_new = lensEquationSolver.findBrightImage(x_source, y_source, kwargs_lens_new, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) print(x_pos, x_pos_new) x_pos = np.sort(x_pos) x_pos_new = np.sort(x_pos_new) y_pos = np.sort(y_pos) y_pos_new = np.sort(y_pos_new) for i in range(len(x_pos)): npt.assert_almost_equal(x_pos[i], x_pos_new[i], decimal=6) npt.assert_almost_equal(y_pos[i], y_pos_new[i], decimal=6) npt.assert_almost_equal(kwargs_lens_new[1]['e1'], kwargs_lens[1]['e1'], decimal=8) npt.assert_almost_equal(kwargs_lens_new[1]['e2'], kwargs_lens[1]['e2'], decimal=8) npt.assert_almost_equal(kwargs_lens_new[0]['e1'], kwargs_lens[0]['e1'], decimal=8) npt.assert_almost_equal(kwargs_lens_new[0]['e2'], kwargs_lens[0]['e2'], decimal=8)
def test_arcs_at_image_position(self): # lensing quantities kwargs_spp = { 'theta_E': 1.26, 'gamma': 2., 'e1': 0.1, 'e2': -0.1, 'center_x': 0.0, 'center_y': 0.0 } # parameters of the deflector lens model # the lens model is a supperposition of an elliptical lens model with external shear lens_model_list = ['SPEP'] #, 'SHEAR'] kwargs_lens = [kwargs_spp] #, kwargs_shear] lens_model_class = LensModel(lens_model_list=lens_model_list) lensEquationSolver = LensEquationSolver(lens_model_class) source_x = 0. source_y = 0.05 x_image, y_image = lensEquationSolver.findBrightImage( source_x, source_y, kwargs_lens, numImages=4, min_distance=0.05, search_window=5) arc_model = LensModel(lens_model_list=['CURVED_ARC_SPP', 'SHIFT']) for i in range(len(x_image)): x0, y0 = x_image[i], y_image[i] print(x0, y0, i) ext = LensModelExtensions(lensModel=lens_model_class) kwargs_arc_i = ext.curved_arc_estimate(x0, y0, kwargs_lens) alpha_x, alpha_y = lens_model_class.alpha(x0, y0, kwargs_lens) kwargs_arc = [ kwargs_arc_i, { 'alpha_x': alpha_x, 'alpha_y': alpha_y } ] print(kwargs_arc_i) direction = kwargs_arc_i['direction'] print(np.cos(direction), np.sin(direction)) x, y = util.make_grid(numPix=5, deltapix=0.01) x = x0 y = y0 gamma1_arc, gamma2_arc = arc_model.gamma(x, y, kwargs_arc) gamma1, gamma2 = lens_model_class.gamma(x, y, kwargs_lens) print(gamma1, gamma2) npt.assert_almost_equal(gamma1_arc, gamma1, decimal=3) npt.assert_almost_equal(gamma2_arc, gamma2, decimal=3) theta_E_arc, gamma_arc, center_x_spp_arc, center_y_spp_arc = CurvedArcSPP.stretch2spp( **kwargs_arc_i) print(theta_E_arc, gamma_arc, center_x_spp_arc, center_y_spp_arc) npt.assert_almost_equal(center_x_spp_arc, 0, decimal=3) npt.assert_almost_equal(center_y_spp_arc, 0, decimal=3)
def test_rayshooting(self): solver = LensEquationSolver(self.lensModel) source_x, source_y = -0.05, -0.02 x_image_true, y_image_true = solver.findBrightImage(source_x, source_y, self.kwargs_epl) fast_rayshooting = MultiplaneFast(x_image_true, y_image_true, self.zlens, self.zsource, self.lens_model_list_epl, self.zlist_epl, astropy_instance=None, param_class=self.param_class, foreground_rays=None, tol_source=1e-5, numerical_alpha_class=None) x_fore, y_fore, alpha_x_fore, alpha_y_fore = fast_rayshooting._ray_shooting_fast_foreground() xtrue, ytrue, alpha_xtrue, alpha_ytrue = self.lensModel.lens_model.\ ray_shooting_partial(np.zeros_like(x_image_true), np.zeros_like(y_image_true), x_image_true, y_image_true, 0., self.zlens, self.kwargs_epl) npt.assert_almost_equal(x_fore, xtrue) npt.assert_almost_equal(y_fore, ytrue) args_lens = self.param_class.kwargs_to_args(self.kwargs_epl) xfast, yfast = fast_rayshooting.ray_shooting_fast(args_lens) x, y = self.lensModel.ray_shooting(x_image_true, y_image_true, self.kwargs_epl) npt.assert_almost_equal(xfast, x) npt.assert_almost_equal(yfast, y) x_inner, y_inner = fast_rayshooting.lensModel.ray_shooting(x_image_true, y_image_true, self.kwargs_epl) npt.assert_almost_equal(x_inner, xfast) npt.assert_almost_equal(y_inner, yfast) foreground_rays = fast_rayshooting._foreground_rays fast_rayshooting_new = MultiplaneFast(x_image_true, y_image_true, self.zlens, self.zsource, self.lens_model_list_epl, self.zlist_epl, astropy_instance=None, param_class=self.param_class, foreground_rays=foreground_rays, tol_source=1e-5, numerical_alpha_class=None) xfast, yfast = fast_rayshooting_new.ray_shooting_fast(args_lens) npt.assert_almost_equal(xfast, x) npt.assert_almost_equal(yfast, y) chi_square_source = fast_rayshooting.source_plane_chi_square(args_lens) chi_square_total = fast_rayshooting.chi_square(args_lens) logL = fast_rayshooting.logL(args_lens) logL_true = -0.5 * chi_square_total npt.assert_almost_equal(logL, logL_true) npt.assert_almost_equal(chi_square_total, chi_square_source)
def test_solver_multiplane(self): lens_model_list = ['SPEP', 'SHEAR', 'SIS'] lensModel = LensModel(lens_model_list, z_source=1, lens_redshift_list=[0.5, 0.5, 0.3], multi_plane=True) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.01 sourcePos_y = -0.01 deltapix = 0.05 numPix = 150 gamma = 1.96 e1, e2 = 0.01, 0.01 kwargs_shear = {'e1': e1, 'e2': e2} # gamma_ext: shear strength, psi_ext: shear angel (in radian) kwargs_spemd = {'theta_E': 1., 'gamma': gamma, 'center_x': 0, 'center_y': 0, 'e1': 0.2, 'e2': 0.03} kwargs_sis = {'theta_E': .1, 'center_x': 1, 'center_y': 0} kwargs_lens = [kwargs_spemd, kwargs_shear, kwargs_sis] x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) print(x_pos, y_pos, 'test positions') kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': 0.1, 'e2': 0, 'center_x': 0., 'center_y': 0}, {'e1': e1, 'e2': e2}, {'theta_E': .1, 'center_x': 1, 'center_y': 0}] solver = Solver4Point(lensModel, solver_type='PROFILE') kwargs_lens_new, accuracy = solver.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init) print(kwargs_lens_new, 'kwargs_lens_new') assert accuracy < 10**(-10) x_source, y_source = lensModel.ray_shooting(x_pos, y_pos, kwargs_lens_new) x_source, y_source = np.mean(x_source), np.mean(y_source) x_pos_new, y_pos_new = lensEquationSolver.findBrightImage(x_source, y_source, kwargs_lens_new, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) print(x_pos, x_pos_new) x_pos = np.sort(x_pos) x_pos_new = np.sort(x_pos_new) y_pos = np.sort(y_pos) y_pos_new = np.sort(y_pos_new) for i in range(len(x_pos)): npt.assert_almost_equal(x_pos[i], x_pos_new[i], decimal=6) npt.assert_almost_equal(y_pos[i], y_pos_new[i], decimal=6) npt.assert_almost_equal(kwargs_lens_new[1]['e1'], kwargs_lens[1]['e1'], decimal=8)
def test_solver_simplified_2(self): lens_model_list = ['SPEP', 'SHEAR_GAMMA_PSI'] lensModel = LensModel(lens_model_list) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = -0.1 deltapix = 0.05 numPix = 150 gamma = 1.96 e1, e2 = -0.01, -0.01 psi_ext, gamma_ext = param_util.ellipticity2phi_gamma(e1, e2) kwargs_shear = {'gamma_ext': gamma_ext, 'psi_ext': psi_ext} # gamma_ext: shear strength, psi_ext: shear angel (in radian) kwargs_spemd = {'theta_E': 1., 'gamma': gamma, 'center_x': 0, 'center_y': 0, 'e1': -0.2, 'e2': -0.03} kwargs_lens = [kwargs_spemd, kwargs_shear] x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': 0, 'e2': 0, 'center_x': 0., 'center_y': 0}, {'gamma_ext': gamma_ext, 'psi_ext': psi_ext}] solver = Solver4Point(lensModel, solver_type='PROFILE_SHEAR') kwargs_lens_new, accuracy = solver.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init) assert accuracy < 10**(-10) x_source, y_source = lensModel.ray_shooting(x_pos, y_pos, kwargs_lens_new) x_source, y_source = np.mean(x_source), np.mean(y_source) x_pos_new, y_pos_new = lensEquationSolver.findBrightImage(x_source, y_source, kwargs_lens_new, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) print(x_pos, x_pos_new) x_pos = np.sort(x_pos) x_pos_new = np.sort(x_pos_new) y_pos = np.sort(y_pos) y_pos_new = np.sort(y_pos_new) for i in range(len(x_pos)): npt.assert_almost_equal(x_pos[i], x_pos_new[i], decimal=6) npt.assert_almost_equal(y_pos[i], y_pos_new[i], decimal=6) npt.assert_almost_equal(kwargs_lens_new[1]['psi_ext'], kwargs_lens[1]['psi_ext'], decimal=8) npt.assert_almost_equal(kwargs_lens_new[1]['gamma_ext'], kwargs_lens[1]['gamma_ext'], decimal=8)
def get_ps_mcmc_kwargs(self, lens_init_kwargs): """Get the init kwargs for point source LENSED_POSITION based on the init kwargs for the lens model, a centered point source, and reasonable redshifts Returns ------- dict kwargs for LENSED_POSITION that can be used to initialize the MCMC """ lens_model_class = LensModel(lens_model_list=['PEMD', 'SHEAR'], z_lens=0.5, z_source=1.5) kwargs_lens = lens_init_kwargs lensEquationSolver = LensEquationSolver(lens_model_class) x_image, y_image = lensEquationSolver.findBrightImage( 0.0, 0.0, # center the point source kwargs_lens, numImages=4, min_distance=0.05, search_window=5.12) magnification = lens_model_class.magnification(x_image, y_image, kwargs=kwargs_lens) fixed_ps = [{}] kwargs_ps_init = [{ 'ra_image': x_image, 'dec_image': y_image, 'point_amp': np.abs(magnification) * 1000 }] kwargs_ps_sigma = [{ 'ra_image': 0.01 * np.ones(len(x_image)), 'dec_image': 0.01 * np.ones(len(x_image)) }] kwargs_lower_ps = [{ 'ra_image': -10 * np.ones(len(x_image)), 'dec_image': -10 * np.ones(len(y_image)) }] kwargs_upper_ps = [{ 'ra_image': 10 * np.ones(len(x_image)), 'dec_image': 10 * np.ones(len(y_image)) }] ps_params = [ kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps ] return ps_params
def test_solver_shapelets(self): lens_model_list = ['SHAPELETS_CART', 'SPEP'] lensModel = LensModel(lens_model_list) solver = Solver4Point(lensModel) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = -0. deltapix = 0.05 numPix = 150 coeffs = np.array([0, 0.1, 0.1, 0, 0, -0.1]) kwargs_lens = [{ 'beta': 1., 'coeffs': coeffs, 'center_x': 0., 'center_y': 0. }, { 'theta_E': 1., 'gamma': 2, 'q': 0.8, 'phi_G': 0.5, 'center_x': 0, 'center_y': 0 }] x_pos, y_pos = lensEquationSolver.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) print(x_pos, y_pos) kwargs_lens_init = [{ 'beta': 1, 'coeffs': np.zeros_like(coeffs), 'center_x': 0., 'center_y': 0 }, kwargs_lens[1]] kwargs_lens_new, accuracy = solver.constraint_lensmodel( x_pos, y_pos, kwargs_lens_init) npt.assert_almost_equal(kwargs_lens_new[0]['beta'], kwargs_lens[0]['beta'], decimal=3) coeffs_new = kwargs_lens_new[0]['coeffs'] for i in range(len(coeffs)): npt.assert_almost_equal(coeffs_new[i], coeffs[i], decimal=3)
def generate_lens(sigma_bkg=sigma_bkg, exp_time=exp_time, numPix=numPix, deltaPix=deltaPix, fwhm=fwhm, psf_type=psf_type, kernel_size=kernel_size, z_source=z_source, z_lens=z_lens, phi_ext=phi_ext, gamma_ext=gamma_ext, theta_E=theta_E, gamma_lens=gamma_lens, e1_lens=e1_lens, e2_lens=e2_lens, center_x_lens_light=center_x_lens_light, center_y_lens_light=center_y_lens_light, source_x=source_y, source_y=source_y, q_source=q_source, phi_source=phi_source, center_x=center_x, center_y=center_y, amp_source=amp_source, R_sersic_source=R_sersic_source, n_sersic_source=n_sersic_source, phi_lens_light=phi_lens_light, q_lens_light=q_lens_light, amp_lens=amp_lens, R_sersic_lens=R_sersic_lens, n_sersic_lens=n_sersic_lens, amp_ps=amp_ps, supersampling_factor=supersampling_factor, v_min=v_min, v_max=v_max, cosmo=cosmo, cosmo2=cosmo2, lens_pos_eq_lens_light_pos=True, same_cosmology=True): kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg) data_class = ImageData(**kwargs_data) kwargs_psf = {'psf_type': psf_type, 'pixel_size': deltaPix, 'fwhm': fwhm} psf_class = PSF(**kwargs_psf) cosmo = FlatLambdaCDM(H0=75, Om0=0.3, Ob0=0.) cosmo = FlatLambdaCDM(H0=65, Om0=0.3, Ob0=0.) gamma1, gamma2 = param_util.shear_polar2cartesian(phi=phi_ext, gamma=gamma_ext) kwargs_shear = {'gamma1': gamma1, 'gamma2': gamma2} if lens_pos_eq_lens_light_pos: center_x = center_x_lens_light center_y = center_y_lens_light if same_cosmology: cosmo2 = cosmo kwargs_spemd = { 'theta_E': theta_E, 'gamma': gamma_lens, 'center_x': center_x, 'center_y': center_y, 'e1': e1_lens, 'e2': e2_lens } lens_model_list = ['SPEP', 'SHEAR'] kwargs_lens = [kwargs_spemd, kwargs_shear] lens_model_class = LensModel(lens_model_list=lens_model_list, z_lens=z_lens, z_source=z_source, cosmo=cosmo) lens_model_class2 = LensModel(lens_model_list=lens_model_list, z_lens=z_lens, z_source=z_source, cosmo=cosmo2) e1_source, e2_source = param_util.phi_q2_ellipticity(phi_source, q_source) kwargs_sersic_source = { 'amp': amp_source, 'R_sersic': R_sersic_source, 'n_sersic': n_sersic_source, 'e1': e1_source, 'e2': e2_source, 'center_x': source_x, 'center_y': source_y } source_model_list = ['SERSIC_ELLIPSE'] kwargs_source = [kwargs_sersic_source] source_model_class = LightModel(light_model_list=source_model_list) ## e1_lens_light, e2_lens_light = param_util.phi_q2_ellipticity( phi_lens_light, q_lens_light) kwargs_sersic_lens = { 'amp': amp_lens, 'R_sersic': R_sersic_lens, 'n_sersic': n_sersic_lens, 'e1': e1_lens_light, 'e2': e2_lens_light, 'center_x': center_x_lens_light, 'center_y': center_y_lens_light } lens_light_model_list = ['SERSIC_ELLIPSE'] kwargs_lens_light = [kwargs_sersic_lens] lens_light_model_class = LightModel(light_model_list=lens_light_model_list) ## lensEquationSolver = LensEquationSolver(lens_model_class) x_image, y_image = lensEquationSolver.findBrightImage( source_x, source_y, #position of ps kwargs_lens, #lens proporties numImages=4, #expected number of images min_distance=deltaPix, #'resolution' search_window=numPix * deltaPix) #search window limits mag = lens_model_class.magnification( x_image, y_image, #for found above ps positions kwargs=kwargs_lens) # and same lens properties kwargs_ps = [{ 'ra_image': x_image, 'dec_image': y_image, 'point_amp': np.abs(mag) * amp_ps }] point_source_list = ['LENSED_POSITION'] point_source_class = PointSource(point_source_type_list=point_source_list, fixed_magnification_list=[False]) kwargs_numerics = {'supersampling_factor': supersampling_factor} imageModel = ImageModel( data_class, # take generated above data specs psf_class, # same for psf lens_model_class, # lens model (gal+ext) source_model_class, # sourse light model lens_light_model_class, # lens light model point_source_class, # add generated ps images kwargs_numerics=kwargs_numerics) image_sim = imageModel.image(kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps) poisson = image_util.add_poisson(image_sim, exp_time=exp_time) bkg = image_util.add_background(image_sim, sigma_bkd=sigma_bkg) image_sim = image_sim + bkg + poisson data_class.update_data(image_sim) kwargs_data['image_data'] = image_sim cmap_string = 'gray' cmap = plt.get_cmap(cmap_string) cmap.set_bad(color='b', alpha=1.) cmap.set_under('k') f, axes = plt.subplots(1, 1, figsize=(6, 6), sharex=False, sharey=False) ax = axes im = ax.matshow(np.log10(image_sim), origin='lower', vmin=v_min, vmax=v_max, cmap=cmap, extent=[0, 1, 0, 1]) ax.get_xaxis().set_visible(False) ax.get_yaxis().set_visible(False) ax.autoscale(False) plt.show() kwargs_model = { 'lens_model_list': lens_model_list, 'lens_light_model_list': lens_light_model_list, 'source_light_model_list': source_model_list, 'point_source_model_list': point_source_list } from lenstronomy.Analysis.td_cosmography import TDCosmography td_cosmo = TDCosmography(z_lens, z_source, kwargs_model, cosmo_fiducial=cosmo) # time delays, the unit [days] is matched when the lensing angles are in arcsec t_days = td_cosmo.time_delays(kwargs_lens, kwargs_ps, kappa_ext=0) dt_days = t_days[1:] - t_days[0] dt_sigma = [3, 5, 10] # Gaussian errors dt_measured = np.random.normal(dt_days, dt_sigma) print("the measured relative delays are: ", dt_measured) return f, source_model_list, kwargs_data, kwargs_psf, kwargs_numerics, dt_measured, dt_sigma, kwargs_ps, lens_model_list, lens_light_model_list, source_model_list, point_source_list, lens_model_class2
def test_all_nfw(self): lensModel = LensModel(['SPEP']) solver_nfw_ellipse = Solver2Point(lensModel, solver_type='ELLIPSE') solver_nfw_center = Solver2Point(lensModel, solver_type='CENTER') spep = LensModel(['SPEP']) image_position_nfw = LensEquationSolver(LensModel(['SPEP', 'NFW'])) sourcePos_x = 0.1 sourcePos_y = 0.03 deltapix = 0.05 numPix = 100 gamma = 1.9 Rs = 0.1 nfw = NFW() alpha_Rs = nfw._rho02alpha(1., Rs) phi_G, q = 0.5, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi_G, q) kwargs_lens = [{ 'theta_E': 1., 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0.1, 'center_y': -0.1 }, { 'Rs': Rs, 'alpha_Rs': alpha_Rs, 'center_x': -0.5, 'center_y': 0.5 }] x_pos, y_pos = image_position_nfw.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=2, min_distance=deltapix, search_window=numPix * deltapix) print(len(x_pos), 'number of images') x_pos = x_pos[:2] y_pos = y_pos[:2] kwargs_init = [{ 'theta_E': 1, 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0., 'center_y': 0 }, { 'Rs': Rs, 'alpha_Rs': alpha_Rs, 'center_x': -0.5, 'center_y': 0.5 }] kwargs_out_center, precision = solver_nfw_center.constraint_lensmodel( x_pos, y_pos, kwargs_init) source_x, source_y = spep.ray_shooting(x_pos[0], y_pos[0], kwargs_out_center) x_pos_new, y_pos_new = image_position_nfw.findBrightImage( source_x, source_y, kwargs_out_center, numImages=2, min_distance=deltapix, search_window=numPix * deltapix) print(kwargs_out_center, 'kwargs_out_center') npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=2) npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=2) npt.assert_almost_equal(kwargs_out_center[0]['center_x'], kwargs_lens[0]['center_x'], decimal=2) npt.assert_almost_equal(kwargs_out_center[0]['center_y'], kwargs_lens[0]['center_y'], decimal=2) npt.assert_almost_equal(kwargs_out_center[0]['center_y'], -0.1, decimal=2) kwargs_init = [{ 'theta_E': 1., 'gamma': gamma, 'e1': 0, 'e2': 0, 'center_x': 0.1, 'center_y': -0.1 }, { 'Rs': Rs, 'alpha_Rs': alpha_Rs, 'center_x': -0.5, 'center_y': 0.5 }] kwargs_out_ellipse, precision = solver_nfw_ellipse.constraint_lensmodel( x_pos, y_pos, kwargs_init) npt.assert_almost_equal(kwargs_out_ellipse[0]['e1'], kwargs_lens[0]['e1'], decimal=2) npt.assert_almost_equal(kwargs_out_ellipse[0]['e2'], kwargs_lens[0]['e2'], decimal=2) npt.assert_almost_equal(kwargs_out_ellipse[0]['e1'], e1, decimal=2)
def test_theta_E_ellipse(self): lensModel = LensModel(['SPEP', 'SHEAR']) solver = Solver2Point(lensModel, solver_type='THETA_E_ELLIPSE') image_position = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = 0.03 deltapix = 0.05 numPix = 100 gamma = 1.9 kwargs_lens = [{ 'theta_E': 1., 'gamma': gamma, 'e1': 0.1, 'e2': -0.03, 'center_x': 0.1, 'center_y': -0.1 }, { 'gamma1': 0.03, 'gamma2': 0.0 }] x_pos, y_pos = image_position.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=2, min_distance=deltapix, search_window=numPix * deltapix, precision_limit=10**(-15)) print(len(x_pos), 'number of images') x_pos = x_pos[:2] y_pos = y_pos[:2] kwargs_init = [{ 'theta_E': 1.9, 'gamma': gamma, 'e1': -0.03, 'e2': 0.1, 'center_x': 0.1, 'center_y': -0.1 }, { 'gamma1': 0.03, 'gamma2': 0.0 }] kwargs_out, precision = solver.constraint_lensmodel( x_pos, y_pos, kwargs_init) print(kwargs_out, 'output') source_x, source_y = lensModel.ray_shooting(x_pos[0], y_pos[0], kwargs_out) x_pos_new, y_pos_new = image_position.findBrightImage( source_x, source_y, kwargs_out, numImages=2, min_distance=deltapix, search_window=numPix * deltapix) npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=3) npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=3) npt.assert_almost_equal(kwargs_out[0]['theta_E'], kwargs_lens[0]['theta_E'], decimal=3) npt.assert_almost_equal(kwargs_out[0]['e1'], kwargs_lens[0]['e1'], decimal=2) npt.assert_almost_equal(kwargs_out[0]['e2'], kwargs_lens[0]['e2'], decimal=2)
def test_magnification_finite_adaptive(self): lens_model_list = ['EPL', 'SHEAR'] z_source = 1.5 kwargs_lens = [{ 'theta_E': 1., 'gamma': 2., 'e1': 0.02, 'e2': -0.09, 'center_x': 0, 'center_y': 0 }, { 'gamma1': 0.01, 'gamma2': 0.03 }] lensmodel = LensModel(lens_model_list) extension = LensModelExtensions(lensmodel) solver = LensEquationSolver(lensmodel) source_x, source_y = 0.07, 0.03 x_image, y_image = solver.findBrightImage(source_x, source_y, kwargs_lens) source_fwhm_parsec = 40. pc_per_arcsec = 1000 / self.cosmo.arcsec_per_kpc_proper(z_source).value source_sigma = source_fwhm_parsec / pc_per_arcsec / 2.355 grid_size = auto_raytracing_grid_size(source_fwhm_parsec) grid_resolution = auto_raytracing_grid_resolution(source_fwhm_parsec) # make this even higher resolution to show convergence grid_number = int(2 * grid_size / grid_resolution) window_size = 2 * grid_size mag_square_grid = extension.magnification_finite( x_image, y_image, kwargs_lens, source_sigma=source_sigma, grid_number=grid_number, window_size=window_size) flux_ratios_square_grid = mag_square_grid / max(mag_square_grid) mag_adaptive_grid = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=self.cosmo) flux_ratios_adaptive_grid = mag_adaptive_grid / max(mag_adaptive_grid) mag_adaptive_grid_2 = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=self.cosmo, axis_ratio=0) flux_ratios_adaptive_grid_2 = mag_adaptive_grid_2 / max( mag_adaptive_grid_2) # tests the default cosmology _ = extension.magnification_finite_adaptive(x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=None, tol=0.0001) # test smallest eigenvalue _ = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=None, tol=0.0001, use_largest_eigenvalue=False) # tests the r_max > sqrt(2) * grid_radius stop criterion _ = extension.magnification_finite_adaptive(x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=None, tol=0.0001, step_size=1000) mag_point_source = abs( lensmodel.magnification(x_image, y_image, kwargs_lens)) quarter_precent_precision = [0.0025] * 4 npt.assert_array_less( flux_ratios_square_grid / flux_ratios_adaptive_grid - 1, quarter_precent_precision) npt.assert_array_less( flux_ratios_square_grid / flux_ratios_adaptive_grid_2 - 1, quarter_precent_precision) half_percent_precision = [0.005] * 4 npt.assert_array_less(mag_square_grid / mag_adaptive_grid - 1, half_percent_precision) npt.assert_array_less(mag_square_grid / mag_adaptive_grid_2 - 1, half_percent_precision) npt.assert_array_less(mag_adaptive_grid / mag_point_source - 1, half_percent_precision) flux_array = np.array([0., 0.]) x_image, y_image = [x_image[0]], [y_image[0]] grid_x = np.array([0., source_sigma]) grid_y = np.array([0., 0.]) grid_r = np.hypot(grid_x, grid_y) source_model = LightModel(['GAUSSIAN']) kwargs_source = [{ 'amp': 1., 'center_x': source_x, 'center_y': source_y, 'sigma': source_sigma }] r_min = 0. r_max = source_sigma * 0.9 flux_array = extension._magnification_adaptive_iteration( flux_array, x_image, y_image, grid_x, grid_y, grid_r, r_min, r_max, lensmodel, kwargs_lens, source_model, kwargs_source) bx, by = lensmodel.ray_shooting(x_image[0], y_image[0], kwargs_lens) sb_true = source_model.surface_brightness(bx, by, kwargs_source) npt.assert_equal(True, flux_array[0] == sb_true) npt.assert_equal(True, flux_array[1] == 0.) r_min = source_sigma * 0.9 r_max = 2 * source_sigma flux_array = extension._magnification_adaptive_iteration( flux_array, x_image, y_image, grid_x, grid_y, grid_r, r_min, r_max, lensmodel, kwargs_lens, source_model, kwargs_source) bx, by = lensmodel.ray_shooting(x_image[0] + source_sigma, y_image[0], kwargs_lens) sb_true = source_model.surface_brightness(bx, by, kwargs_source) npt.assert_equal(True, flux_array[1] == sb_true)
def test_all_spep(self): lensModel = LensModel(['SPEP']) solver_spep_center = Solver2Point(lensModel, solver_type='CENTER') solver_spep_ellipse = Solver2Point(lensModel, solver_type='ELLIPSE') image_position_spep = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = 0.03 gamma = 1.9 phi_G, q = 0.5, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi_G, q) kwargs_lens = [{ 'theta_E': 1, 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0.1, 'center_y': -0.1 }] x_pos, y_pos = image_position_spep.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=2, min_distance=0.01, search_window=5, precision_limit=10**(-10), num_iter_max=10) print(x_pos, y_pos, 'test') x_pos = x_pos[:2] y_pos = y_pos[:2] kwargs_init = [{ 'theta_E': 1, 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0, 'center_y': 0 }] kwargs_out_center, precision = solver_spep_center.constraint_lensmodel( x_pos, y_pos, kwargs_init) kwargs_init = [{ 'theta_E': 1, 'gamma': gamma, 'e1': 0, 'e2': 0, 'center_x': 0.1, 'center_y': -0.1 }] kwargs_out_ellipse, precision = solver_spep_ellipse.constraint_lensmodel( x_pos, y_pos, kwargs_init) npt.assert_almost_equal(kwargs_out_center[0]['center_x'], kwargs_lens[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_out_center[0]['center_y'], kwargs_lens[0]['center_y'], decimal=3) npt.assert_almost_equal(kwargs_out_center[0]['center_y'], -0.1, decimal=3) npt.assert_almost_equal(kwargs_out_ellipse[0]['e1'], kwargs_lens[0]['e1'], decimal=3) npt.assert_almost_equal(kwargs_out_ellipse[0]['e2'], kwargs_lens[0]['e2'], decimal=3) npt.assert_almost_equal(kwargs_out_ellipse[0]['e1'], e1, decimal=3)
def test_all_spep_sis(self): lensModel = LensModel(['SPEP', 'SIS']) solver_ellipse = Solver2Point(lensModel, solver_type='ELLIPSE') solver_center = Solver2Point(lensModel, solver_type='CENTER') spep = LensModel(['SPEP', 'SIS']) image_position = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = 0.03 deltapix = 0.05 numPix = 100 gamma = 1.9 kwargs_lens = [{ 'theta_E': 1., 'gamma': gamma, 'e1': 0.2, 'e2': -0.03, 'center_x': 0.1, 'center_y': -0.1 }, { 'theta_E': 0.6, 'center_x': -0.5, 'center_y': 0.5 }] x_pos, y_pos = image_position.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=2, min_distance=deltapix, search_window=numPix * deltapix, precision_limit=10**(-10)) print(len(x_pos), 'number of images') x_pos = x_pos[:2] y_pos = y_pos[:2] kwargs_init = [{ 'theta_E': 1, 'gamma': gamma, 'e1': 0.2, 'e2': -0.03, 'center_x': 0., 'center_y': 0 }, { 'theta_E': 0.6, 'center_x': -0.5, 'center_y': 0.5 }] kwargs_out_center, precision = solver_center.constraint_lensmodel( x_pos, y_pos, kwargs_init) print(kwargs_out_center, 'output') source_x, source_y = spep.ray_shooting(x_pos[0], y_pos[0], kwargs_out_center) x_pos_new, y_pos_new = image_position.findBrightImage( source_x, source_y, kwargs_out_center, numImages=2, min_distance=deltapix, search_window=numPix * deltapix) npt.assert_almost_equal(x_pos_new[0], x_pos[0], decimal=3) npt.assert_almost_equal(y_pos_new[0], y_pos[0], decimal=3) npt.assert_almost_equal(kwargs_out_center[0]['center_x'], kwargs_lens[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_out_center[0]['center_y'], kwargs_lens[0]['center_y'], decimal=3) npt.assert_almost_equal(kwargs_out_center[0]['center_y'], -0.1, decimal=3) kwargs_init = [{ 'theta_E': 1., 'gamma': gamma, 'e1': 0, 'e2': 0, 'center_x': 0.1, 'center_y': -0.1 }, { 'theta_E': 0.6, 'center_x': -0.5, 'center_y': 0.5 }] kwargs_out_ellipse, precision = solver_ellipse.constraint_lensmodel( x_pos, y_pos, kwargs_init) npt.assert_almost_equal(kwargs_out_ellipse[0]['e1'], kwargs_lens[0]['e1'], decimal=3) npt.assert_almost_equal(kwargs_out_ellipse[0]['e2'], kwargs_lens[0]['e2'], decimal=3) npt.assert_almost_equal(kwargs_out_ellipse[0]['e1'], 0.2, decimal=3)
#============================================================================== from lenstronomy.PointSource.point_source import PointSource from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver np.random.seed(seed) amp_qRh_s_plane=np.random.uniform(0.8,1./0.8) qso_amp= 10.**(-0.4*(source_para['mag_sersic']-zp))*amp_qRh_s_plane # add_qso = int(input("add QSO?:\n input 0 no qso, others add qso:\t")) add_qso= 1 # add_qso= 0 if add_qso == 0: qso_amp = 0 lensEquationSolver = LensEquationSolver(lens_model_class) x_image, y_image = lensEquationSolver.findBrightImage(source_pos[0], source_pos[1], kwargs_lens_list, numImages=4, min_distance=deltaPix, search_window=numPix * deltaPix) mag = lens_model_class.magnification(x_image, y_image, kwargs=kwargs_lens_list) kwargs_ps = {'ra_image': x_image, 'dec_image': y_image, 'point_amp': np.abs(mag) * qso_amp} # quasar point source position in the source plane and intrinsic brightness point_source_list = ['LENSED_POSITION'] point_source_class = PointSource(point_source_type_list=point_source_list, fixed_magnification_list=[False]) imageModel_without_arc = ImageModel(data_class, psf_class, lens_model_class=lens_model_class, source_model_class=None, # No arc, i.e. source_model_class=None lens_light_model_class= lens_light_model_class, point_source_class= point_source_class, kwargs_numerics=kwargs_numerics) image_without_arc = imageModel_without_arc.image(kwargs_lens = kwargs_lens_list, kwargs_source = [{}], kwargs_lens_light = kwargs_lens_light_list, kwargs_ps = [kwargs_ps]) image_deflector = imageModel_without_arc.image(kwargs_lens = kwargs_lens_list, kwargs_source = [{}], kwargs_lens_light = kwargs_lens_light_list, kwargs_ps = [kwargs_ps], point_source_add=False, unconvolved=True) imageModel_source = ImageModel(data_class, psf_class, lens_model_class=None,
class Optimizer(object): """ class which executes the optimization routines. Currently implemented as a particle swarm optimization followed by a downhill simplex routine. Particle swarm optimizer is modified from the CosmoHammer particle swarm routine with different convergence criteria implemented. """ def __init__(self, x_pos, y_pos, redshift_list=[], lens_model_list=[], kwargs_lens=[], optimizer_routine='fixed_powerlaw_shear', magnification_target=None, multiplane=None, z_main=None, z_source=None, tol_source=1e-5, tol_mag=0.2, tol_centroid=0.05, centroid_0=[0, 0], astropy_instance=None, verbose=False, re_optimize=False, particle_swarm=True, pso_convergence_standardDEV=0.01, pso_convergence_mean=10000, pso_compute_magnification=500, tol_simplex_params=1e-3, tol_simplex_func=1e-3, tol_src_penalty=0.1, constrain_params=None, simplex_n_iterations=400, compute_mags_postpso=False, optimizer_kwargs={}): """ :param x_pos: observed position in arcsec :param y_pos: observed position in arcsec :param magnification_target: observed magnifications, uncertainty for the magnifications :param redshift_list: list of lens model redshifts :param lens_model_list: list of lens models :param kwargs_lens: keywords for lens models :param optimizer_routine: a set optimization routine; currently only 'optimize_SIE_shear' is implemented :param multiplane: multi-plane flag :param z_main: if multi-plane, macromodel redshift :param z_source: if multi-plane, macromodel redshift :param tol_source: tolerance on the source position in the optimization :param tol_mag: tolerance on the magnification :param tol_centroid: tolerance on the centroid of the mass distributions :param centroid_0: centroid of the optimized lens model (list [x,y]) :param astropy_instance: instance of astropy :param interpolate: if multi-plane; flag to interpolate the background lens models :param verbose: flag to print status updates during optimization :param re_optimize: flag to run in re-optimization mode; the particle swarm particles will be initialized clustered around the values in kwargs_lens :param particle_swarm: if False, the optimizer will skip the particle swarm optimization and go straight to downhill simplex :param pso_convergence_standardDEV: convergence criterion for particle swarm algorithm :param pso_convergence_mean: alternate convergence criterion for PSO; usually dominates the former :param pso_compute_magnification: flag for computing magnifications in the PSO; useful to avoid computing the magnification for lens models that are obviously wrong :param tol_simplex_func/params: tolerance for the scipy downhill simplex routine :param tol_src_penalty: if the source penalty is less than this value, the recomputation of the image positions will be skipped. (if the source penalty is good enough, you're guaranteed to match the input image positions so not point in recomputing them) :param constrain_params: additional parameters to constrain (type: dictionary) Format: {'parameter_name_1':[desired_value,uncertainty], 'parameter_name_2':[desired_value,uncertainty]} e.g. {'theta_E:[1,0.01]} will constrain the Einstein radius to 1 plus/minus 0.01 The parameter name must be part of the 'params_to_vary' attribute of the specific optimization routine used (see class 'fixed_routines') Special cases are constraining shear and shear_pa in polar coordinates: {'shear':[0.05,0.01], 'shear_pa':[30,5]} will constrain the shear parameters in polar coordinates based on the cartesian e1/e2 values :param simplex_n_iterations: simplex_n_iterations times problem dimension gives the maximum # of iterations for the downhill simplex routine :param optimizer_kwargs: optional keyword arguments for the mutliplane optimizer :param compute_mags_postpso: flag to automatically compute magnifications when perfomring downhill simplex optimization. Note: if running with particle_swarm = False, the re_optimize variable does nothing """ self._multiplane = multiplane self._verbose = verbose x_pos, y_pos = np.array(x_pos), np.array(y_pos) self.x_pos, self.y_pos = x_pos, y_pos self._re_optimize = re_optimize self._particle_swarm = particle_swarm self._init_kwargs = kwargs_lens self._pso_convergence_standardDEV = pso_convergence_standardDEV self._tol_simplex_params = tol_simplex_params self._tol_simplex_func = tol_simplex_func self._tol_src_penalty = tol_src_penalty self._simplex_iter = simplex_n_iterations self._compute_mags_postpso = compute_mags_postpso if 'optimization_algorithm' in optimizer_kwargs: self._opt_method = optimizer_kwargs['optimization_algorithm'] else: self._opt_method = 'Nelder-Mead' # make sure the length of observed positions matches, length of observed magnifications, etc. self._init_test(x_pos, y_pos, magnification_target, tol_source, redshift_list, lens_model_list, kwargs_lens, z_source, z_main, multiplane, astropy_instance) # initialize lens model class self._lensModel = LensModel(lens_model_list=lens_model_list, redshift_list=redshift_list, z_source=z_source, cosmo=astropy_instance, multi_plane=multiplane) # initiate a params class that, based on the optimization routine, determines which parameters/lens models to optimize self._params = Params(zlist=self._lensModel.redshift_list, lens_list=self._lensModel.lens_model_list, arg_list=kwargs_lens, optimizer_routine=optimizer_routine, xpos=x_pos, ypos=y_pos) # initialize particle swarm inital param limits if 're_optimize_scale' in optimizer_kwargs: scale = optimizer_kwargs['re_optimize_scale'] self._re_optimize_scale = scale else: scale = 1 self._re_optimize_scale = scale self._lower_limit, self._upper_limit = self._params.to_vary_limits( self._re_optimize, scale=scale) # initiate optimizer classes, one for particle swarm and one for the downhill simplex if multiplane is False: lensing_class = SinglePlaneLensing(self._lensModel, x_pos, y_pos, self._params, kwargs_lens) # don't bother with anything special here, just use the regular lensmodel class self.solver = LensEquationSolver(self._lensModel) else: lensing_class = MultiPlaneLensing(self._lensModel, x_pos, y_pos, kwargs_lens, z_source, z_main, astropy_instance, self._params.tovary_indicies, optimizer_kwargs) self.solver = LensEquationSolver(lensing_class) self.lensModel = self.solver.lensModel self._optimizer = Penalties( tol_source, tol_mag, tol_centroid, lensing_class, centroid_0, magnification_target, params_to_constrain=constrain_params, param_class=self._params, pso_convergence_mean=pso_convergence_mean, pso_compute_magnification=pso_compute_magnification, compute_mags=False, verbose=verbose) def optimize(self, n_particles=50, n_iterations=250, restart=1): """ :param n_particles: number of particle swarm particles :param n_iterations: number of particle swarm iternations :param restart: number of times to execute the optimization; the best result of all optimizations will be returned. total number of lens models sovled: n_particles*n_iterations :return: lens model keywords, [optimized source position], best fit image positions """ if restart < 0: raise ValueError( "parameter 'restart' must be integer of value > 0") # particle swarm optimization penalties, parameters, src_pen_best = [], [], [] for run in range(0, restart): penalty, params = self._single_optimization( n_particles, n_iterations) penalties.append(penalty) parameters.append(params) src_pen_best.append(self._optimizer.src_pen_best) # select the best optimization best_index = np.argmin(penalties) # combine the optimized parameters with the parameters kept fixed during the optimization to obtain full kwargs_lens kwargs_varied = self._params.argstovary_todictionary( parameters[best_index]) kwargs_lens_final = kwargs_varied + self._params.argsfixed_todictionary( ) # solve for the optimized image positions srcx, srcy = self._optimizer.lensing._ray_shooting_fast(kwargs_varied) source_x, source_y = np.mean(srcx), np.mean(srcy) # if we have a good enough solution, no point in recomputing the image positions since this can be quite slow # and will give the same answer if src_pen_best[best_index] < self._tol_src_penalty: x_image, y_image = self.x_pos, self.y_pos else: # Here, the solver has the instance of "lensing_class" or "LensModel" for multiplane/singleplane respectively. print('Warning: possibly a bad fit.') x_image, y_image = self.solver.findBrightImage( source_x, source_y, kwargs_lens_final, arrival_time_sort=False) #x_image, y_image = self.solver.image_position_from_source(source_x, source_y, kwargs_lens_final, arrival_time_sort = False) if self._verbose: print('optimization done.') print('Recovered source position: ', (srcx, srcy)) return kwargs_lens_final, [source_x, source_y], [x_image, y_image] def _single_optimization(self, n_particles, n_iterations): self._optimizer._reset() self._optimizer._init_particles(n_particles, n_iterations) if self._particle_swarm: params = self._pso(n_particles, n_iterations, self._optimizer) if self._verbose: print('PSO done.') else: params = self._params._kwargs_to_tovary(self._init_kwargs) if self._verbose: print('starting amoeba... ') # downhill simplex optimization self._optimizer._reset(compute_mags=self._compute_mags_postpso) if self._opt_method == 'Nelder-Mead' or self._opt_method == 'powell': options = { 'adaptive': True, 'fatol': self._tol_simplex_func, 'xatol': self._tol_simplex_params, 'maxiter': self._simplex_iter * len(params) } else: options = {'maxiter': self._simplex_iter * len(params)} optimized_downhill_simplex = minimize(self._optimizer, x0=params, method=self._opt_method, options=options) penalty = self._optimizer._get_best() self._optimizer._reset() return penalty, optimized_downhill_simplex['x'] def _pso(self, n_particles, n_iterations, optimizer): """ :param n_particles: number of PSO particles :param n_iterations: number of PSO iterations :param optimizer: instance of SinglePlaneOptimizer or MultiPlaneOptimizer :return: optimized kwargs_lens """ pso = ParticleSwarmOptimizer(optimizer, low=self._lower_limit, high=self._upper_limit, particleCount=n_particles) gBests = pso._optimize(maxIter=n_iterations, standard_dev=self._pso_convergence_standardDEV) likelihoods = [particle.fitness for particle in gBests] ind = np.argmax(likelihoods) return gBests[ind].position def _init_test(self, x_pos, y_pos, magnification_target, tol_source, zlist, lens_list, arg_list, z_source, z_main, multiplane, astropy_instance): """ check inputs """ assert len(x_pos) == len(y_pos) assert len(magnification_target) == len(x_pos) assert tol_source is not None assert len(lens_list) == len(arg_list) if multiplane is True: assert len(zlist) == len(lens_list) assert z_source is not None assert z_main is not None assert astropy_instance is not None
def setup(self): # data specifics sigma_bkg = .05 # background noise per pixel (Gaussian) exp_time = 100. # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit) numPix = 100 # cutout pixel size deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2) fwhm = 0.1 # full width half max of PSF (only valid when psf_type='gaussian') psf_type = 'GAUSSIAN' # 'GAUSSIAN', 'PIXEL', 'NONE' # generate the coordinate grid and image properties kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg) kwargs_data['exposure_time'] = exp_time * np.ones_like(kwargs_data['image_data']) data_class = ImageData(**kwargs_data) # generate the psf variables kwargs_psf = {'psf_type': psf_type, 'pixel_size': deltaPix, 'fwhm': fwhm} # kwargs_psf = sim_util.psf_configure_simple(psf_type=psf_type, fwhm=fwhm, kernelsize=kernel_size, deltaPix=deltaPix, kernel=kernel) psf_class = PSF(**kwargs_psf) # lensing quantities kwargs_shear = {'gamma1': 0.02, 'gamma2': -0.04} # shear values to the source plane kwargs_spemd = {'theta_E': 1.26, 'gamma': 2., 'center_x': 0.0, 'center_y': 0.0, 'e1': -0.1, 'e2': 0.05} # parameters of the deflector lens model # the lens model is a supperposition of an elliptical lens model with external shear lens_model_list = ['EPL', 'SHEAR'] kwargs_lens_true = [kwargs_spemd, kwargs_shear] lens_model_class = LensModel(lens_model_list=lens_model_list) # choice of source type source_type = 'SERSIC' # 'SERSIC' or 'SHAPELETS' source_x = 0. source_y = 0.05 # Sersic parameters in the initial simulation phi_G, q = 0.5, 0.8 e1, e2 = param_util.phi_q2_ellipticity(phi_G, q) kwargs_sersic_source = {'amp': 1000, 'R_sersic': 0.05, 'n_sersic': 1, 'e1': e1, 'e2': e2, 'center_x': source_x, 'center_y': source_y} # kwargs_else = {'sourcePos_x': source_x, 'sourcePos_y': source_y, 'quasar_amp': 400., 'gamma1_foreground': 0.0, 'gamma2_foreground':-0.0} source_model_list = ['SERSIC_ELLIPSE'] kwargs_source_true = [kwargs_sersic_source] source_model_class = LightModel(light_model_list=source_model_list) lensEquationSolver = LensEquationSolver(lens_model_class) x_image, y_image = lensEquationSolver.findBrightImage(source_x, source_y, kwargs_lens_true, numImages=4, min_distance=deltaPix, search_window=numPix * deltaPix) mag = lens_model_class.magnification(x_image, y_image, kwargs=kwargs_lens_true) kwargs_numerics = {'supersampling_factor': 1} imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, kwargs_numerics=kwargs_numerics) # generate image model = imageModel.image(kwargs_lens_true, kwargs_source_true) poisson = image_util.add_poisson(model, exp_time=exp_time) bkg = image_util.add_background(model, sigma_bkd=sigma_bkg) image_sim = model + bkg + poisson data_class.update_data(image_sim) kwargs_data['image_data'] = image_sim kwargs_model = {'lens_model_list': lens_model_list, 'source_light_model_list': source_model_list, } # make cutous and data instances of them x_pos, y_pos = data_class.map_coord2pix(x_image, y_image) ra_grid, dec_grid = data_class.pixel_coordinates multi_band_list = [] for i in range(len(x_pos)): n_cut = 12 x_c = int(x_pos[i]) y_c = int(y_pos[i]) image_cut = image_sim[int(y_c - n_cut):int(y_c + n_cut), int(x_c - n_cut):int(x_c + n_cut)] exposure_map_cut = data_class.exposure_map[int(y_c - n_cut):int(y_c + n_cut), int(x_c - n_cut):int(x_c + n_cut)] kwargs_data_i = { 'background_rms': data_class.background_rms, 'exposure_time': exposure_map_cut, 'ra_at_xy_0': ra_grid[y_c - n_cut, x_c - n_cut], 'dec_at_xy_0': dec_grid[y_c - n_cut, x_c - n_cut], 'transform_pix2angle': data_class.transform_pix2angle , 'image_data': image_cut } multi_band_list.append([kwargs_data_i, kwargs_psf, kwargs_numerics]) kwargs_params = {'kwargs_lens': kwargs_lens_true, 'kwargs_source': kwargs_source_true} self.multiPatch = MultiPatchPlot(multi_band_list, kwargs_model, kwargs_params, multi_band_type='joint-linear', kwargs_likelihood=None, verbose=True, cmap_string="gist_heat") self.data_class = data_class self.model = model self.lens_model_class = lens_model_class self.kwargs_lens = kwargs_lens_true
def main(): args = parse_args() cfg = Config.fromfile(args.config) if args.n_data is not None: cfg.n_data = args.n_data # Seed for reproducibility np.random.seed(cfg.seed) random.seed(cfg.seed) if not os.path.exists(cfg.out_dir): os.makedirs(cfg.out_dir) print("Destination folder: {:s}".format(cfg.out_dir)) else: raise OSError("Destination folder already exists.") # Instantiate PSF models psf_models = get_PSF_models(cfg.psf, cfg.instrument.pixel_scale) n_psf = len(psf_models) # Instantiate ImageData #kwargs_data = sim_util.data_configure_simple(**cfg.image) #image_data = ImageData(**kwargs_data) # Instantiate density models kwargs_model = dict( lens_model_list=[ cfg.bnn_omega.lens_mass.profile, cfg.bnn_omega.external_shear.profile ], source_light_model_list=[cfg.bnn_omega.src_light.profile], ) lens_mass_model = LensModel( lens_model_list=kwargs_model['lens_model_list']) src_light_model = LightModel( light_model_list=kwargs_model['source_light_model_list']) lens_eq_solver = LensEquationSolver(lens_mass_model) lens_light_model = None ps_model = None if 'lens_light' in cfg.components: kwargs_model['lens_light_model_list'] = [ cfg.bnn_omega.lens_light.profile ] lens_light_model = LightModel( light_model_list=kwargs_model['lens_light_model_list']) if 'agn_light' in cfg.components: kwargs_model['point_source_model_list'] = [ cfg.bnn_omega.agn_light.profile ] ps_model = PointSource( point_source_type_list=kwargs_model['point_source_model_list'], fixed_magnification_list=[False]) # Initialize BNN prior bnn_prior = getattr(bnn_priors, cfg.bnn_prior_class)(cfg.bnn_omega, cfg.components) # Initialize dataframe of labels param_list = [] for comp in cfg.components: param_list += [ '{:s}_{:s}'.format(comp, param) for param in bnn_prior.params[cfg.bnn_omega[comp]['profile']] ] if 'agn_light' in cfg.components: param_list += ['magnification_{:d}'.format(i) for i in range(4)] param_list += ['x_image_{:d}'.format(i) for i in range(4)] param_list += ['y_image_{:d}'.format(i) for i in range(4)] param_list += ['n_img'] param_list += ['img_path', 'total_magnification'] if cfg.bnn_prior_class == 'EmpiricalBNNPrior': param_list += [ 'z_lens', 'z_src', 'vel_disp_iso', 'R_eff_lens', 'R_eff_src', 'abmag_src' ] metadata = pd.DataFrame(columns=param_list) print("Starting simulation...") current_idx = 0 # running idx of dataset pbar = tqdm(total=cfg.n_data) while current_idx < cfg.n_data: psf_model = psf_models[current_idx % n_psf] sample = bnn_prior.sample() # FIXME: sampling in batches if sample['lens_mass']['theta_E'] < cfg.selection.theta_E.min: continue # Instantiate SimAPI (converts mag to amp and wraps around image model) kwargs_detector = util.merge_dicts(cfg.instrument, cfg.bandpass, cfg.observation) kwargs_detector.update( seeing=cfg.psf.fwhm, psf_type=cfg.psf.type, kernel_point_source=psf_model ) # keyword deprecation warning: I asked Simon to change this to data_api = DataAPI(cfg.image.num_pix, **kwargs_detector) image_data = data_api.data_class #sim_api = SimAPI(numpix=cfg.image.num_pix, # kwargs_single_band=kwargs_detector, # kwargs_model=kwargs_model, # kwargs_numerics=cfg.numerics) kwargs_lens_mass = [sample['lens_mass'], sample['external_shear']] kwargs_src_light = [sample['src_light']] kwargs_src_light = amp_to_mag_extended(kwargs_src_light, src_light_model, data_api) kwargs_lens_light = None kwargs_ps = None if 'agn_light' in cfg.components: x_image, y_image = lens_eq_solver.findBrightImage( sample['src_light']['center_x'], sample['src_light']['center_y'], kwargs_lens_mass, numImages=4, min_distance=cfg.instrument.pixel_scale, search_window=cfg.image.num_pix * cfg.instrument.pixel_scale) magnification = np.abs( lens_mass_model.magnification(x_image, y_image, kwargs=kwargs_lens_mass)) unlensed_mag = sample['agn_light']['magnitude'] # unlensed agn mag kwargs_unlensed_mag_ps = [{ 'ra_image': x_image, 'dec_image': y_image, 'magnitude': unlensed_mag }] # note unlensed magnitude kwargs_unlensed_amp_ps = amp_to_mag_point( kwargs_unlensed_mag_ps, ps_model, data_api) # note unlensed amp kwargs_ps = copy.deepcopy(kwargs_unlensed_amp_ps) for kw in kwargs_ps: kw.update(point_amp=kw['point_amp'] * magnification) else: kwargs_unlensed_amp_ps = None if 'lens_light' in cfg.components: kwargs_lens_light = [sample['lens_light']] kwargs_lens_light = amp_to_mag_extended(kwargs_lens_light, lens_light_model, data_api) # Instantiate image model image_model = ImageModel(image_data, psf_model, lens_mass_model, src_light_model, lens_light_model, ps_model, kwargs_numerics=cfg.numerics) # Compute magnification lensed_total_flux = get_lensed_total_flux(kwargs_lens_mass, kwargs_src_light, kwargs_lens_light, kwargs_ps, image_model) unlensed_total_flux = get_unlensed_total_flux(kwargs_src_light, src_light_model, kwargs_unlensed_amp_ps, ps_model) total_magnification = lensed_total_flux / unlensed_total_flux # Apply magnification cut if total_magnification < cfg.selection.magnification.min: continue # Generate image for export img = image_model.image(kwargs_lens_mass, kwargs_src_light, kwargs_lens_light, kwargs_ps) #kwargs_in_amp = sim_api.magnitude2amplitude(kwargs_lens_mass, kwargs_src_light, kwargs_lens_light, kwargs_ps) #imsim_api = sim_api.image_model_class #imsim_api.image(*kwargs_in_amp) # Add noise noise = data_api.noise_for_model(img, background_noise=True, poisson_noise=True, seed=None) img += noise # Save image file img_path = os.path.join(cfg.out_dir, 'X_{0:07d}.npy'.format(current_idx + 1)) np.save(img_path, img) # Save labels meta = {} for comp in cfg.components: for param_name, param_value in sample[comp].items(): meta['{:s}_{:s}'.format(comp, param_name)] = param_value if 'agn_light' in cfg.components: n_img = len(x_image) for i in range(n_img): meta['magnification_{:d}'.format(i)] = magnification[i] meta['x_image_{:d}'.format(i)] = x_image[i] meta['y_image_{:d}'.format(i)] = y_image[i] meta['n_img'] = n_img if cfg.bnn_prior_class == 'EmpiricalBNNPrior': for misc_name, misc_value in sample['misc'].items(): meta['{:s}'.format(misc_name)] = misc_value meta['total_magnification'] = total_magnification meta['img_path'] = img_path metadata = metadata.append(meta, ignore_index=True) # Update progress current_idx += 1 pbar.update(1) pbar.close() # Fix column ordering metadata = metadata[param_list] metadata_path = os.path.join(cfg.out_dir, 'metadata.csv') metadata.to_csv(metadata_path, index=None) print("Labels include: ", metadata.columns.values)
def sim_lens(data, numPix=101, sigma_bkg=8.0, exp_time=100.0, deltaPix=0.263, psf_type='GAUSSIAN', kernel_size=91): flux_g = mag_to_flux(data['mag_g'], 27.5) flux_r = mag_to_flux(data['mag_r'], 27.5) flux_i = mag_to_flux(data['mag_i'], 27.5) flux_z = mag_to_flux(data['mag_z'], 27.5) flux_source = mag_to_flux(data['source_mag'], 27.5) flux_lens = mag_to_flux(data['lens_mag'], 27.5) color_idx = {'g': 0, 'r': 1, 'i': 2, 'z': 3} cosmo = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.) full_band_images = np.zeros((numPix, numPix, 4)) ### Set kwargs based on input data file #shear kwargs_shear = { 'gamma_ext': data['lens_shear_gamma_ext'], 'psi_ext': data['lens_shear_psi_ext'] } #lens potential kwargs_spemd = { 'theta_E': data['lens_theta_E'], 'gamma': data['lens_gamma'], 'center_x': data['lens_center_x'], 'center_y': data['lens_center_y'], 'e1': data['lens_e1'], 'e2': data['lens_e2'] } #lens light kwargs_sersic_lens = { 'amp': flux_lens, 'R_sersic': data['lens_R_sersic'], 'n_sersic': data['lens_n_sersic'], 'e1': data['lens_e1'], 'e2': data['lens_e2'], 'center_x': data['lens_center_x'], 'center_y': data['lens_center_y'] } #source kwargs_sersic_source = { 'amp': flux_source, 'R_sersic': data['source_R_sersic'], 'n_sersic': data['source_n_sersic'], 'e1': data['source_e1'], 'e2': data['source_e2'], 'center_x': data['source_center_x'], 'center_y': data['source_center_y'] } ###set model parameters based on kwargs #lens potential lens_model_list = ['SPEP', 'SHEAR_GAMMA_PSI'] kwargs_lens = [kwargs_spemd, kwargs_shear] lens_model_class = LensModel(lens_model_list=lens_model_list) #lens light lens_light_model_list = ['SERSIC_ELLIPSE'] kwargs_lens_light = [kwargs_sersic_lens] lens_light_model_class = LightModel(light_model_list=lens_light_model_list) #source source_model_list = ['SERSIC_ELLIPSE'] kwargs_source = [kwargs_sersic_source] source_model_class = LightModel(light_model_list=source_model_list) ###configure image based on data properties kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, sigma_bkg) data_class = ImageData(**kwargs_data) ###solve lens equation lensEquationSolver = LensEquationSolver(lens_model_class) x_image, y_image = lensEquationSolver.findBrightImage( kwargs_sersic_source['center_x'], kwargs_sersic_source['center_y'], kwargs_lens, numImages=4, min_distance=deltaPix, search_window=numPix * deltaPix) magnification = lens_model_class.magnification(x_image, y_image, kwargs=kwargs_lens) ###iterate through bands to simulate images for band in ['g', 'r', 'i', 'z']: #psf info kwargs_psf = { 'psf_type': psf_type, 'fwhm': data['psf_%s' % band], 'pixel_size': deltaPix, 'truncation': 3 } psf_class = PSF(**kwargs_psf) #quasar info kwargs_ps = [{ 'ra_image': x_image, 'dec_image': y_image, 'point_amp': np.abs(magnification) * eval('flux_%s' % band) }] point_source_list = ['LENSED_POSITION'] point_source_class = PointSource( point_source_type_list=point_source_list, fixed_magnification_list=[False]) #build image model kwargs_numerics = {'supersampling_factor': 1} imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics) #generate image image_sim = imageModel.image(kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps) poisson = image_util.add_poisson(image_sim, exp_time=exp_time) bkg = image_util.add_background(image_sim, sigma_bkd=sigma_bkg) image_sim = image_sim + bkg + poisson data_class.update_data(image_sim) kwargs_data['image_data'] = image_sim kwargs_model = { 'lens_model_list': lens_model_list, 'lens_light_model_list': lens_light_model_list, 'source_light_model_list': source_model_list, 'point_source_model_list': point_source_list } #build up an array with one slice for each band full_band_images[:, :, color_idx[band]] += image_sim return full_band_images
def test_decoupling(self): lens_model_list = ['SPEP', 'SIS'] lensModel = LensModel(lens_model_list) solver = Solver4Point(lensModel, decoupling=False) solver_decoupled = Solver4Point(lensModel, decoupling=True) lensEquationSolver = LensEquationSolver(lensModel) sourcePos_x = 0.1 sourcePos_y = -0.1 deltapix = 0.05 numPix = 150 gamma = 1.9 kwargs_lens = [{ 'theta_E': 1., 'gamma': gamma, 'q': 0.8, 'phi_G': 0.5, 'center_x': 0.1, 'center_y': -0.1 }, { 'theta_E': 0.1, 'center_x': 0.5, 'center_y': 0 }] x_pos, y_pos = lensEquationSolver.findBrightImage( sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix * deltapix) kwargs_lens_init = [{ 'theta_E': 1.3, 'gamma': gamma, 'q': 0.9, 'phi_G': 1.5, 'center_x': 0., 'center_y': 0 }, { 'theta_E': 0.1, 'center_x': 0.5, 'center_y': 0 }] kwargs_lens_new, accuracy = solver.constraint_lensmodel( x_pos, y_pos, kwargs_lens_init) kwargs_lens_new_2, accuracy = solver_decoupled.constraint_lensmodel( x_pos, y_pos, kwargs_lens_init) print(kwargs_lens_new_2) print(kwargs_lens_new) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], kwargs_lens[0]['theta_E'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['q'], kwargs_lens[0]['q'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['phi_G'], kwargs_lens[0]['phi_G'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_x'], kwargs_lens[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_y'], kwargs_lens[0]['center_y'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], kwargs_lens_new_2[0]['theta_E'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['q'], kwargs_lens_new_2[0]['q'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['phi_G'], kwargs_lens_new_2[0]['phi_G'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_x'], kwargs_lens_new_2[0]['center_x'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['center_y'], kwargs_lens_new_2[0]['center_y'], decimal=3) npt.assert_almost_equal(kwargs_lens_new[0]['theta_E'], 1., decimal=3) lensModel = LensModel(lens_model_list=lens_model_list) x_source_new, y_source_new = lensModel.ray_shooting( x_pos, y_pos, kwargs_lens_new) dist = np.sqrt((x_source_new - x_source_new[0])**2 + (y_source_new - y_source_new[0])**2) print(dist) assert np.max(dist) < 0.000001
def main(): args = parse_args() input_dir = args.datadir object_type = args.object_type # Load DB files as dataframes lens_df = pd.read_sql(f'{object_type}_lens', str('sqlite:///' + os.path.join(input_dir, 'lens_truth.db')), index_col=0) ps_df = pd.read_sql( f'lensed_{object_type}', str('sqlite:///' + os.path.join(input_dir, f'lensed_{object_type}_truth.db')), index_col=0) src_light_df = pd.read_sql(f'{object_type}_hosts', str('sqlite:///' + os.path.join(input_dir, 'host_truth.db')), index_col=0) # Init columns to add ps_df['total_magnification'] = np.nan src_light_df['total_magnification_bulge'] = np.nan src_light_df['total_magnification_disk'] = np.nan for band in list('ugrizy'): src_light_df[f'lensed_flux_{band}'] = np.nan src_light_df[f'lensed_flux_{band}_noMW'] = np.nan dc2_sprinkler = DC2Sprinkler() # utility class for flux integration ##################### # Model assumptions # ##################### # Cosmology cosmo = WMAP7 # DC2 #cosmo = wCDM(H0=72.0, Om0=0.26, Ode0=0.74, w0=-1.0) # OM10 # Imaging config, only used to set the numerics of the lens equation solver pixel_scale = args.pixel_size num_pix = args.num_pix # Density profiles kwargs_model = {} kwargs_model['lens_model_list'] = ['SIE', 'SHEAR_GAMMA_PSI'] kwargs_model['point_source_model_list'] = ['LENSED_POSITION'] arcsec_to_deg = 1 / 3600.0 # Instantiate tool for imaging our hosts lensed_host_imager = lensing_utils.LensedHostImager(pixel_scale, num_pix) sys_ids = lens_df['lens_cat_sys_id'].unique() progress = tqdm(total=len(sys_ids)) for i, sys_id in enumerate(sys_ids): ####################### # Slice relevant rows # ####################### # Lens mass lens_info = lens_df[lens_df['lens_cat_sys_id'] == sys_id].iloc[0].squeeze() # Lensed point source ps_info = ps_df[ps_df['lens_cat_sys_id'] == sys_id].iloc[[ 0 ]].copy() # sub-df of length 1, to be updated ps_df.drop( ps_df[ps_df['lens_cat_sys_id'] == sys_id].index, axis=0, inplace=True) # delete rows for this system for now, will add back # Host light src_light_read_only = src_light_df.loc[ src_light_df['lens_cat_sys_id'] == sys_id].iloc[0].squeeze( ) # for accessing the original host info; arbitarily take the first lensed image, since properties we use are same across the images src_light_info = src_light_df[ src_light_df['lens_cat_sys_id'] == sys_id].iloc[[ 0 ]].copy() # sub-df of length 1, to be updated src_light_df.drop( src_light_df[src_light_df['lens_cat_sys_id'] == sys_id].index, axis=0, inplace=True) # delete rows for this system for now, will add back # Properties defining lens geometry z_lens = lens_info['redshift'] z_src = src_light_read_only['redshift'] x_lens = lens_info['ra_lens'] # absolute position in rad y_lens = lens_info['dec_lens'] x_src = ps_info[f'x_{object_type}'].values[ 0] # defined in arcsec wrt lens center y_src = ps_info[f'y_{object_type}'].values[0] x_src_host = src_light_read_only[ 'x_src'] # defined in arcsec wrt lens center y_src_host = src_light_read_only['y_src'] ######################################## # Solve lens equation for point source # ######################################## lens_mass_model = LensModel([ 'SIE', 'SHEAR_GAMMA_PSI', ], cosmo=cosmo, z_lens=z_lens, z_source=z_src) lens_eq_solver = LensEquationSolver(lens_mass_model) kwargs_lens_mass = lensing_utils.get_lens_params(lens_info, z_src=z_src, cosmo=cosmo) x_image, y_image = lens_eq_solver.findBrightImage( x_src, y_src, kwargs_lens_mass, min_distance=0.01, # default is 0.01 numImages=4, search_window=num_pix * pixel_scale, # default is 5 precision_limit=10**(-10) # default ) magnification = np.abs( lens_mass_model.magnification(x_image, y_image, kwargs=kwargs_lens_mass)) td_cosmo = TDCosmography(z_lens, z_src, kwargs_model, cosmo_fiducial=cosmo) ps_kwargs = [{'ra_image': x_image, 'dec_image': y_image}] time_delays = td_cosmo.time_delays(kwargs_lens_mass, ps_kwargs, kappa_ext=0.0) n_img = len(x_image) ################################# # Update lensed AGN truth table # ################################# ps_info = ps_info.loc[ps_info.index.repeat(n_img)].reset_index( drop=True) # replicate enough rows # Absolute image positions in rad ra_image_abs = x_lens + x_image * arcsec_to_deg / np.cos( np.radians(y_lens)) dec_image_abs = y_lens + y_image * arcsec_to_deg # Reorder the existing images by increasing dec to enforce a consistent image ordering system # Only 'unique_id', 'image_number', ra', 'dec', 't_delay', 'magnification' are affected by the reordering increasing_dec_i = np.argsort(dec_image_abs) ps_info['unique_id'] = [ '{:s}_{:d}'.format(ps_info.iloc[0]['dc2_sys_id'], img_i) for img_i in range(n_img) ] ps_info['image_number'] = np.arange(n_img) ps_info['ra'] = ra_image_abs[increasing_dec_i] ps_info['dec'] = dec_image_abs[increasing_dec_i] ps_info['magnification'] = magnification[increasing_dec_i] time_delays = time_delays[increasing_dec_i] time_delays -= time_delays[0] # time delays relative to first image ps_info['t_delay'] = time_delays #ps_df.update(ps_info) # inplace op doesn't work when n_img is different from OM10 # FIXME: Check that the following works to update fluxes correctly if object_type == 'agn': # since AGN follow a single SED template agn_magnorm = ps_info['magnorm'].iloc[ 0] # unlensed magnorm, same across images for agn_idx, agn_magnorm in list( enumerate(ps_info['magnorm'].values)): dmag = -2.5 * np.log10( np.abs(ps_info['magnification'].iloc[agn_idx])) magnorm_dict = { band: agn_magnorm + dmag for band in ['u', 'g', 'r', 'i', 'z', 'y'] } agn_flux_no_mw, agn_flux_mw = dc2_sprinkler.add_flux( 'agnSED/agn.spec.gz', z_src, magnorm_dict, lens_info['av_mw'], lens_info['rv_mw']) for band in list('ugrizy'): ps_info.loc[agn_idx, f'flux_{band}_agn'] = agn_flux_mw[band] ps_info.loc[agn_idx, f'flux_{band}_agn_noMW'] = agn_flux_no_mw[band] ps_info['total_magnification'] = np.sum(np.abs(magnification)) ps_df = ps_df.append(ps_info, ignore_index=True, sort=False) ps_df.reset_index(drop=True, inplace=True) # to prevent duplicate indices ######################################### # Solve lens equation for host centroid # ######################################### # Solve the lens equation for a hypothetical point source at the host centroid x_image_host, y_image_host = lens_eq_solver.findBrightImage( x_src_host, y_src_host, kwargs_lens_mass, min_distance=0.0075, # default is 0.01 numImages=4, search_window=num_pix * pixel_scale, # default is 5 precision_limit=10**(-10) # default ) # Absolute image positions in rad ra_image_abs_host = x_lens + x_image_host * arcsec_to_deg / np.cos( np.radians(y_lens)) dec_image_abs_host = y_lens + y_image_host * arcsec_to_deg n_img_host = len(y_image_host) ########################### # Update host truth table # ########################### src_light_info = src_light_info.loc[src_light_info.index.repeat( n_img_host)].reset_index(drop=True) # replicate enough rows # Reorder the existing images by increasing dec to enforce a consistent image ordering system increasing_dec_i_host = np.argsort(y_image_host) src_light_info['unique_id'] = [ '{:s}_{:d}'.format(src_light_read_only['dc2_sys_id'], img_i) for img_i in range(n_img_host) ] src_light_info['image_number'] = np.arange(n_img_host) src_light_info['ra_host_lensed'] = ra_image_abs_host[ increasing_dec_i_host] src_light_info['dec_host_lensed'] = dec_image_abs_host[ increasing_dec_i_host] src_light_info['x_img'] = x_image_host[increasing_dec_i_host] src_light_info['y_img'] = y_image_host[increasing_dec_i_host] bulge_img, bulge_features = lensed_host_imager.get_image( lens_info, src_light_read_only, z_lens, z_src, 'bulge') disk_img, disk_features = lensed_host_imager.get_image( lens_info, src_light_read_only, z_lens, z_src, 'disk') # Update magnorms based on lensed and unlensed flux for band in list('ugrizy'): src_light_info[f'magnorm_bulge_{band}'] = bulge_features[ 'magnorms'][band] src_light_info[f'magnorm_disk_{band}'] = disk_features['magnorms'][ band] src_light_info['total_magnification_bulge'] = bulge_features[ 'total_magnification'] src_light_info['total_magnification_disk'] = disk_features[ 'total_magnification'] disk_flux_no_mw, disk_flux_mw = dc2_sprinkler.add_flux( src_light_read_only['sed_disk_host'][2:-1], z_src, disk_features['magnorms'], lens_info['av_mw'], lens_info['rv_mw']) bulge_flux_no_mw, bulge_flux_mw = dc2_sprinkler.add_flux( src_light_read_only['sed_bulge_host'][2:-1], z_src, bulge_features['magnorms'], lens_info['av_mw'], lens_info['rv_mw']) for band in list('ugrizy'): src_light_info[f'lensed_flux_{band}'] = disk_flux_mw[ band] + bulge_flux_mw[band] src_light_info[f'lensed_flux_{band}_noMW'] = disk_flux_no_mw[ band] + bulge_flux_no_mw[band] src_light_df = src_light_df.append(src_light_info, ignore_index=True, sort=False) src_light_df.reset_index(drop=True, inplace=True) # to prevent duplicate indices progress.update(1) # Sort by dc2_sys_id and image number ps_df['dc2_sys_id_int'] = [ int(sys_id.split('_')[-1]) for sys_id in ps_df['dc2_sys_id'] ] ps_df.sort_values(['dc2_sys_id_int', 'image_number'], axis=0, inplace=True) ps_df.drop(['dc2_sys_id_int'], axis=1, inplace=True) src_light_df['dc2_sys_id_int'] = [ int(sys_id.split('_')[-1]) for sys_id in src_light_df['dc2_sys_id'] ] src_light_df.sort_values(['dc2_sys_id_int', 'image_number'], axis=0, inplace=True) src_light_df.drop(['dc2_sys_id_int'], axis=1, inplace=True) # Export lensed_ps and host truth tables to original file format io_utils.export_db(ps_df, input_dir, f'updated_lensed_{object_type}_truth.db', f'lensed_{object_type}', overwrite=True) io_utils.export_db(src_light_df, input_dir, f'updated_host_truth.db', f'{object_type}_hosts', overwrite=True) progress.close()
class Imager: """Deterministic utility class for imaging the objects on a pixel grid Attributes ---------- bnn_omega : dict copy of `cfg.bnn_omega` components : list list of components, e.g. `lens_mass` """ def __init__(self, components, lens_mass_model, src_light_model, lens_light_model=None, ps_model=None, kwargs_numerics={'supersampling_factor': 1}, min_magnification=0.0, for_cosmography=False, magnification_frac_err=0.0): self.components = components self.kwargs_numerics = kwargs_numerics self.lens_mass_model = lens_mass_model self.src_light_model = src_light_model self.lens_light_model = lens_light_model self.ps_model = ps_model self.unlensed_ps_model = PointSource(point_source_type_list=['SOURCE_POSITION'], fixed_magnification_list=[False]) self.lens_eq_solver = LensEquationSolver(self.lens_mass_model) self.min_magnification = min_magnification self.for_cosmography = for_cosmography self.magnification_frac_err = magnification_frac_err self.img_features = {} # Initialized to store metadata of images, will get updated for each lens def _set_sim_api(self, num_pix, kwargs_detector, psf_kernel_size, which_psf_maps): """Set the simulation API objects """ self.data_api = DataAPI(num_pix, **kwargs_detector) #self.pixel_scale = data_api.pixel_scale pixel_scale = kwargs_detector['pixel_scale'] psf_model = psf_utils.get_PSF_model(kwargs_detector['psf_type'], pixel_scale, seeing=kwargs_detector['seeing'], kernel_size=psf_kernel_size, which_psf_maps=which_psf_maps) # Set the precision level of lens equation solver self.min_distance = 0.05 self.search_window = pixel_scale*num_pix self.image_model = ImageModel(self.data_api.data_class, psf_model, self.lens_mass_model, self.src_light_model, self.lens_light_model, self.ps_model, kwargs_numerics=self.kwargs_numerics) if 'agn_light' in self.components: self.unlensed_image_model = ImageModel(self.data_api.data_class, psf_model, None, self.src_light_model, None, self.unlensed_ps_model, kwargs_numerics=self.kwargs_numerics) else: self.unlensed_image_model = ImageModel(self.data_api.data_class, psf_model, None, self.src_light_model, None, None, kwargs_numerics=self.kwargs_numerics) def _load_kwargs(self, sample): """Generate an image from provided model and model parameters Parameters ---------- sample : dict model parameters sampled by a bnn_prior object """ self._load_lens_mass_kwargs(sample['lens_mass'], sample['external_shear']) self._load_src_light_kwargs(sample['src_light']) if 'lens_light' in self.components: self._load_lens_light_kwargs(sample['lens_light']) else: self.kwargs_lens_light = None if 'agn_light' in self.components: self._load_agn_light_kwargs(sample) else: self.kwargs_ps = None self.kwargs_unlensed_unmagnified_amp_ps = None def _load_lens_mass_kwargs(self, lens_mass_sample, external_shear_sample): self.kwargs_lens_mass = [lens_mass_sample, external_shear_sample] def _load_src_light_kwargs(self, src_light_sample): kwargs_src_light = [src_light_sample] # Convert from mag to amp self.kwargs_src_light = mag_to_amp_extended(kwargs_src_light, self.src_light_model, self.data_api) def _load_lens_light_kwargs(self, lens_light_sample): kwargs_lens_light = [lens_light_sample] # Convert lens magnitude into amp self.kwargs_lens_light = mag_to_amp_extended(kwargs_lens_light, self.lens_light_model, self.data_api) def _load_agn_light_kwargs(self, sample): """Set the point source kwargs to be ingested by Lenstronomy """ # When using the image positions for cosmological parameter recovery, the time delays must be computed by evaluating the Fermat potential at these exact positions. if self.for_cosmography: x_image = sample['misc']['x_image'] y_image = sample['misc']['y_image'] # When the precision of the lens equation solver doesn't have to be matched between image positions and time delays, simply solve for the image positions using whatever desired precision. else: x_image, y_image = self.lens_eq_solver.findBrightImage(self.kwargs_src_light[0]['center_x'], self.kwargs_src_light[0]['center_y'], self.kwargs_lens_mass, min_distance=self.min_distance, search_window=self.search_window, numImages=4, num_iter_max=100, # default is 10 but td_cosmography default is 100 precision_limit=10**(-10) # default for both this and td_cosmography ) agn_light_sample = sample['agn_light'] unlensed_mag = agn_light_sample['magnitude'] # unlensed agn mag # Save the unlensed (source-plane) kwargs in amplitude units kwargs_unlensed_unmagnified_mag_ps = [{'ra_source': self.kwargs_src_light[0]['center_x'], 'dec_source': self.kwargs_src_light[0]['center_y'], 'magnitude': unlensed_mag}] self.kwargs_unlensed_unmagnified_amp_ps = mag_to_amp_point(kwargs_unlensed_unmagnified_mag_ps, self.unlensed_ps_model, self.data_api) # note # Compute the lensed (image-plane), magnified kwargs in amplitude units magnification = self.lens_mass_model.magnification(x_image, y_image, kwargs=self.kwargs_lens_mass) measured_magnification = np.abs(magnification*(1.0 + self.magnification_frac_err*np.random.randn(len(magnification)))) # Add noise to magnification magnification = np.abs(magnification) kwargs_lensed_unmagnified_mag_ps = [{'ra_image': x_image, 'dec_image': y_image, 'magnitude': unlensed_mag}] # note unlensed magnitude kwargs_lensed_unmagnified_amp_ps = mag_to_amp_point(kwargs_lensed_unmagnified_mag_ps, self.ps_model, self.data_api) # note unmagnified amp self.kwargs_ps = copy.deepcopy(kwargs_lensed_unmagnified_amp_ps) for kw in self.kwargs_ps: kw.update(point_amp=kw['point_amp']*measured_magnification) # Log the solved image positions self.img_features.update(x_image=x_image, y_image=y_image, magnification=magnification, measured_magnification=measured_magnification) def generate_image(self, sample, num_pix, survey_object_dict): img_canvas = np.empty([len(survey_object_dict), num_pix, num_pix]) # [n_filters, num_pix, num_pix] # Loop over bands for i, (bp, survey_object) in enumerate(survey_object_dict.items()): self._set_sim_api(num_pix, survey_object.kwargs_single_band(), survey_object.psf_kernel_size, survey_object.which_psf_maps) self._load_kwargs(sample) # Reject nonsensical number of images (due to insufficient numerical precision) if ('y_image' in self.img_features) and (len(self.img_features['y_image']) not in [2, 4]): return None, None # Compute magnification lensed_total_flux = get_lensed_total_flux(self.kwargs_lens_mass, self.kwargs_src_light, self.kwargs_ps, self.image_model) #unlensed_total_flux = get_unlensed_total_flux(self.kwargs_src_light, self.src_light_model, self.kwargs_unlensed_amp_ps, self.ps_model) unlensed_total_flux = get_unlensed_total_flux_numerical(self.kwargs_src_light, self.kwargs_unlensed_unmagnified_amp_ps, self.unlensed_image_model) total_magnification = lensed_total_flux/unlensed_total_flux # Apply magnification cut if (total_magnification < self.min_magnification) or np.isnan(total_magnification): return None, None # Generate image for export img = self.image_model.image(self.kwargs_lens_mass, self.kwargs_src_light, self.kwargs_lens_light, self.kwargs_ps) img = np.maximum(0.0, img) # safeguard against negative pixel values img_canvas[i, :, :] = img # Save remaining image features img_features_single_band = {f'total_magnification_{bp}': total_magnification, f'lensed_total_flux_{bp}': lensed_total_flux, f'unlensed_total_flux_{bp}': unlensed_total_flux} self.img_features.update(img_features_single_band) return img_canvas, self.img_features def add_noise(self, image_array): """Add noise to the image (deprecated; replaced by the data_augmentation package) """ #noise_map = self.data_api.noise_for_model(image_array, background_noise=True, poisson_noise=True, seed=None) #image_array += noise_map #return image_array pass
def test_magnification_finite_adaptive(self): lens_model_list = ['EPL', 'SHEAR'] z_source = 1.5 kwargs_lens = [{ 'theta_E': 1., 'gamma': 2., 'e1': 0.02, 'e2': -0.09, 'center_x': 0, 'center_y': 0 }, { 'gamma1': 0.01, 'gamma2': 0.03 }] lensmodel = LensModel(lens_model_list) extension = LensModelExtensions(lensmodel) solver = LensEquationSolver(lensmodel) source_x, source_y = 0.07, 0.03 x_image, y_image = solver.findBrightImage(source_x, source_y, kwargs_lens) source_fwhm_parsec = 40. pc_per_arcsec = 1000 / self.cosmo.arcsec_per_kpc_proper(z_source).value source_sigma = source_fwhm_parsec / pc_per_arcsec / 2.355 grid_size = auto_raytracing_grid_size(source_fwhm_parsec) grid_resolution = auto_raytracing_grid_resolution(source_fwhm_parsec) # make this even higher resolution to show convergence grid_number = int(2 * grid_size / grid_resolution) window_size = 2 * grid_size mag_square_grid = extension.magnification_finite( x_image, y_image, kwargs_lens, source_sigma=source_sigma, grid_number=grid_number, window_size=window_size) flux_ratios_square_grid = mag_square_grid / max(mag_square_grid) mag_adaptive_grid = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=self.cosmo) flux_ratios_adaptive_grid = mag_adaptive_grid / max(mag_adaptive_grid) mag_adaptive_grid_fixed_aperture_size = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=self.cosmo, fixed_aperture_size=True, grid_radius_arcsec=0.2) flux_ratios_fixed_aperture_size = mag_adaptive_grid_fixed_aperture_size / max( mag_adaptive_grid_fixed_aperture_size) mag_adaptive_grid_2 = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=self.cosmo, axis_ratio=0) flux_ratios_adaptive_grid_2 = mag_adaptive_grid_2 / max( mag_adaptive_grid_2) # tests the default cosmology _ = extension.magnification_finite_adaptive(x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=None, tol=0.0001) # test smallest eigenvalue _ = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=None, tol=0.0001, use_largest_eigenvalue=False) # tests the r_max > sqrt(2) * grid_radius stop criterion _ = extension.magnification_finite_adaptive(x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=None, tol=0.0001, step_size=1000) mag_point_source = abs( lensmodel.magnification(x_image, y_image, kwargs_lens)) quarter_precent_precision = [0.0025] * 4 npt.assert_array_less( flux_ratios_square_grid / flux_ratios_adaptive_grid - 1, quarter_precent_precision) npt.assert_array_less( flux_ratios_fixed_aperture_size / flux_ratios_adaptive_grid - 1, quarter_precent_precision) npt.assert_array_less( flux_ratios_square_grid / flux_ratios_adaptive_grid_2 - 1, quarter_precent_precision) half_percent_precision = [0.005] * 4 npt.assert_array_less(mag_square_grid / mag_adaptive_grid - 1, half_percent_precision) npt.assert_array_less(mag_square_grid / mag_adaptive_grid_2 - 1, half_percent_precision) npt.assert_array_less(mag_adaptive_grid / mag_point_source - 1, half_percent_precision) flux_array = np.array([0., 0.]) grid_x = np.array([0., source_sigma]) grid_y = np.array([0., 0.]) grid_r = np.hypot(grid_x, grid_y) # test that the double gaussian profile has 2x the flux when dx, dy = 0 magnification_double_gaussian = extension.magnification_finite_adaptive( x_image, y_image, source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=self.cosmo, source_light_model='DOUBLE_GAUSSIAN', dx=0., dy=0., amp_scale=1., size_scale=1.) npt.assert_almost_equal(magnification_double_gaussian, 2 * mag_adaptive_grid) grid_radius = 0.3 npix = 400 _x = _y = np.linspace(-grid_radius, grid_radius, npix) resolution = 2 * grid_radius / npix xx, yy = np.meshgrid(_x, _y) for i in range(0, 4): beta_x, beta_y = lensmodel.ray_shooting(x_image[i] + xx.ravel(), y_image[i] + yy.ravel(), kwargs_lens) source_light_model = LightModel(['GAUSSIAN'] * 2) amp_scale, dx, dy, size_scale = 0.2, 0.005, -0.005, 1. kwargs_source = [{ 'amp': 1., 'center_x': source_x, 'center_y': source_y, 'sigma': source_sigma }, { 'amp': amp_scale, 'center_x': source_x + dx, 'center_y': source_y + dy, 'sigma': source_sigma * size_scale }] sb = source_light_model.surface_brightness(beta_x, beta_y, kwargs_source) magnification_double_gaussian_reference = np.sum( sb) * resolution**2 magnification_double_gaussian = extension.magnification_finite_adaptive( [x_image[i]], [y_image[i]], source_x, source_y, kwargs_lens, source_fwhm_parsec, z_source, cosmo=self.cosmo, source_light_model='DOUBLE_GAUSSIAN', dx=dx, dy=dy, amp_scale=amp_scale, size_scale=size_scale, grid_resolution=resolution, grid_radius_arcsec=grid_radius, axis_ratio=1.) npt.assert_almost_equal( magnification_double_gaussian / magnification_double_gaussian_reference, 1., 3) source_model = LightModel(['GAUSSIAN']) kwargs_source = [{ 'amp': 1., 'center_x': source_x, 'center_y': source_y, 'sigma': source_sigma }] r_min = 0. r_max = source_sigma * 0.9 flux_array = extension._magnification_adaptive_iteration( flux_array, x_image[0], y_image[0], grid_x, grid_y, grid_r, r_min, r_max, lensmodel, kwargs_lens, source_model, kwargs_source) bx, by = lensmodel.ray_shooting(x_image[0], y_image[0], kwargs_lens) sb_true = source_model.surface_brightness(bx, by, kwargs_source) npt.assert_equal(True, flux_array[0] == sb_true) npt.assert_equal(True, flux_array[1] == 0.) r_min = source_sigma * 0.9 r_max = 2 * source_sigma flux_array = extension._magnification_adaptive_iteration( flux_array, x_image[0], y_image[0], grid_x, grid_y, grid_r, r_min, r_max, lensmodel, kwargs_lens, source_model, kwargs_source) bx, by = lensmodel.ray_shooting(x_image[0] + source_sigma, y_image[0], kwargs_lens) sb_true = source_model.surface_brightness(bx, by, kwargs_source) npt.assert_equal(True, flux_array[1] == sb_true)