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)
