class TestLensingOperator(object):
"""
tests the Lensing Operator class
"""
def setup(self):
self.num_pix = 25 # cutout pixel size
delta_pix = 0.24
_, _, ra_at_xy_0, dec_at_xy_0, _, _, Mpix2coord, _ \
= l_util.make_grid_with_coordtransform(numPix=self.num_pix, deltapix=delta_pix, subgrid_res=1,
inverse=False, left_lower=False)
kwargs_data = {
#'background_rms': background_rms,
#'exposure_time': np.ones((self.num_pix, self.num_pix)) * exp_time, # individual exposure time/weight per pixel
'ra_at_xy_0': ra_at_xy_0,
'dec_at_xy_0': dec_at_xy_0,
'transform_pix2angle': Mpix2coord,
'image_data': np.zeros((self.num_pix, self.num_pix))
}
self.data = ImageData(**kwargs_data)
self.lens_model = LensModel(['SPEP'])
self.kwargs_lens = [{
'theta_E': 1,
'gamma': 2,
'center_x': 0,
'center_y': 0,
'e1': -0.05,
'e2': 0.05
}]
self.kwargs_lens_null = [{
'theta_E': 0,
'gamma': 2,
'center_x': 0,
'center_y': 0,
'e1': 0,
'e2': 0
}]
# PSF specification
kwargs_psf = {'psf_type': 'NONE'}
self.psf = PSF(**kwargs_psf)
# list of source light profiles
source_model_list = ['SERSIC_ELLIPSE']
kwargs_sersic_ellipse_source = {
'amp': 2000,
'R_sersic': 0.6,
'n_sersic': 1,
'e1': 0.1,
'e2': 0.1,
'center_x': 0.3,
'center_y': 0.3
}
kwargs_source = [kwargs_sersic_ellipse_source]
source_model = LightModel(light_model_list=source_model_list)
# list of lens light profiles
lens_light_model_list = []
kwargs_lens_light = [{}]
lens_light_model = LightModel(light_model_list=lens_light_model_list)
kwargs_numerics = {
'supersampling_factor': 1,
'supersampling_convolution': False
}
self.image_model = ImageModel(self.data,
self.psf,
self.lens_model,
source_model,
lens_light_model,
point_source_class=None,
kwargs_numerics=kwargs_numerics)
self.image_grid_class = self.image_model.ImageNumerics.grid_class
self.source_grid_class_default = NumericsSubFrame(self.data,
self.psf).grid_class
# create simulated image
image_sim_no_noise = self.image_model.image(self.kwargs_lens,
kwargs_source,
kwargs_lens_light)
self.source_light_lensed = image_sim_no_noise
self.data.update_data(image_sim_no_noise)
# source only, in source plane, on same grid as data
self.source_light_delensed = self.image_model.source_surface_brightness(
kwargs_source, unconvolved=False, de_lensed=True)
# define some auto mask for tests
self.likelihood_mask = np.zeros_like(self.source_light_lensed)
self.likelihood_mask[self.source_light_lensed > 0.1 *
self.source_light_lensed.max()] = 1
def test_matrix_product(self):
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='nearest_legacy')
lensing_op.update_mapping(self.kwargs_lens)
lensing_op_mat = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='nearest')
lensing_op_mat.update_mapping(self.kwargs_lens)
source_1d = util.image2array(self.source_light_delensed)
image_1d = util.image2array(self.source_light_lensed)
npt.assert_equal(lensing_op.source2image(source_1d),
lensing_op_mat.source2image(source_1d))
npt.assert_equal(lensing_op.image2source(image_1d),
lensing_op_mat.image2source(image_1d))
def test_minimal_source_plane(self):
source_1d = util.image2array(self.source_light_delensed)
# test with no mask
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='nearest',
minimal_source_plane=True)
lensing_op.update_mapping(self.kwargs_lens)
image_1d = util.image2array(self.source_light_lensed)
assert lensing_op.image2source(image_1d).size < source_1d.size
# test with mask
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='nearest',
minimal_source_plane=True)
lensing_op.set_likelihood_mask(self.likelihood_mask)
lensing_op.update_mapping(self.kwargs_lens)
image_1d = util.image2array(self.source_light_lensed)
assert lensing_op.image2source(image_1d).size < source_1d.size
# for 'bilinear' operator, only works with no mask (for now)
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='bilinear',
minimal_source_plane=True)
lensing_op.update_mapping(self.kwargs_lens)
image_1d = util.image2array(self.source_light_lensed)
assert lensing_op.image2source(image_1d).size < source_1d.size
def test_legacy_mapping(self):
"""testing than image2source / source2image are close to the parametric mapping"""
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='nearest_legacy')
lensing_op.update_mapping(self.kwargs_lens)
source_1d = util.image2array(self.source_light_delensed)
image_1d = util.image2array(self.source_light_lensed)
source_1d_lensed = lensing_op.source2image(source_1d)
image_1d_delensed = lensing_op.image2source(image_1d)
assert source_1d_lensed.shape == image_1d.shape
assert image_1d_delensed.shape == source_1d.shape
npt.assert_almost_equal(source_1d_lensed / source_1d_lensed.max(),
image_1d / image_1d.max(),
decimal=0.6)
npt.assert_almost_equal(image_1d_delensed / image_1d_delensed.max(),
source_1d / source_1d.max(),
decimal=0.6)
def test_simple_mapping(self):
"""testing than image2source / source2image are close to the parametric mapping"""
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='nearest')
lensing_op.update_mapping(self.kwargs_lens)
source_1d = util.image2array(self.source_light_delensed)
image_1d = util.image2array(self.source_light_lensed)
source_1d_lensed = lensing_op.source2image(source_1d)
image_1d_delensed = lensing_op.image2source(image_1d)
assert source_1d_lensed.shape == image_1d.shape
assert image_1d_delensed.shape == source_1d.shape
npt.assert_almost_equal(source_1d_lensed / source_1d_lensed.max(),
image_1d / image_1d.max(),
decimal=0.6)
npt.assert_almost_equal(image_1d_delensed / image_1d_delensed.max(),
source_1d / source_1d.max(),
decimal=0.6)
def test_interpol_mapping(self):
"""testing than image2source / source2image are close to the parametric mapping"""
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='bilinear')
lensing_op.update_mapping(self.kwargs_lens)
source_1d = util.image2array(self.source_light_delensed)
image_1d = util.image2array(self.source_light_lensed)
source_1d_lensed = lensing_op.source2image(source_1d)
image_1d_delensed = lensing_op.image2source(image_1d)
assert source_1d_lensed.shape == image_1d.shape
assert image_1d_delensed.shape == source_1d.shape
npt.assert_almost_equal(source_1d_lensed / source_1d_lensed.max(),
image_1d / image_1d.max(),
decimal=0.8)
npt.assert_almost_equal(image_1d_delensed / image_1d_delensed.max(),
source_1d / source_1d.max(),
decimal=0.8)
def test_source2image(self):
lensing_op = LensingOperator(self.lens_model, self.image_grid_class,
self.source_grid_class_default,
self.num_pix)
source_1d = util.image2array(self.source_light_delensed)
source_1d_lensed = lensing_op.source2image(
source_1d, kwargs_lens=self.kwargs_lens)
assert len(source_1d_lensed.shape) == 1
source_2d = self.source_light_delensed
source_2d_lensed = lensing_op.source2image_2d(
source_2d, kwargs_lens=self.kwargs_lens, update_mapping=True)
assert len(source_2d_lensed.shape) == 2
def test_image2source(self):
lensing_op = LensingOperator(self.lens_model, self.image_grid_class,
self.source_grid_class_default,
self.num_pix)
image_1d = util.image2array(self.source_light_lensed)
image_1d_delensed = lensing_op.image2source(
image_1d, kwargs_lens=self.kwargs_lens)
assert len(image_1d_delensed.shape) == 1
image_2d = self.source_light_lensed
image_2d_delensed = lensing_op.image2source_2d(
image_2d, kwargs_lens=self.kwargs_lens, update_mapping=True)
assert len(image_2d_delensed.shape) == 2
def test_source_plane_coordinates(self):
lensing_op = LensingOperator(self.lens_model, self.image_grid_class,
self.source_grid_class_default,
self.num_pix)
theta_x, theta_y = lensing_op.source_plane_coordinates
assert theta_x.size == self.num_pix**2
assert theta_y.size == self.num_pix**2
subgrid_res = 2
source_grid_class = NumericsSubFrame(
self.data, self.psf, supersampling_factor=subgrid_res).grid_class
lensing_op = LensingOperator(self.lens_model, self.image_grid_class,
source_grid_class, self.num_pix)
theta_x, theta_y = lensing_op.source_plane_coordinates
assert theta_x.size == self.num_pix**2 * subgrid_res**2
assert theta_y.size == self.num_pix**2 * subgrid_res**2
def test_image_plane_coordinates(self):
lensing_op = LensingOperator(self.lens_model, self.image_grid_class,
self.source_grid_class_default,
self.num_pix)
theta_x, theta_y = lensing_op.image_plane_coordinates
assert theta_x.size == self.num_pix**2
assert theta_y.size == self.num_pix**2
def test_find_source_pixel(self):
lensing_op = LensingOperator(self.lens_model,
self.image_grid_class,
self.source_grid_class_default,
self.num_pix,
source_interpolation='nearest')
beta_x, beta_y = self.lens_model.ray_shooting(
lensing_op.imagePlane.theta_x, lensing_op.imagePlane.theta_y,
self.kwargs_lens)
i = 10
j = lensing_op._find_source_pixel_nearest_legacy(i, beta_x, beta_y)
assert (isinstance(j, int) or isinstance(j, np.int64))
def __init__(self, *args, **kwargs):
super(TestRaise, self).__init__(*args, **kwargs)
# data specifics
sigma_bkg = .05 # background noise per pixel
exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
numPix = 100 # cutout pixel size
deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2)
fwhm = 0.5 # full width half max of PSF
# PSF specification
kwargs_data = sim_util.data_configure_simple(numPix,
deltaPix,
exp_time,
sigma_bkg,
inverse=True)
self.data_class = ImageData(**kwargs_data)
kwargs_psf = {
'psf_type': 'GAUSSIAN',
'fwhm': fwhm,
'truncation': 5,
'pixel_size': deltaPix
}
psf_class = PSF(**kwargs_psf)
kernel = psf_class.kernel_point_source
kwargs_psf = {
'psf_type': 'PIXEL',
'kernel_point_source': kernel,
'psf_error_map': np.ones_like(kernel) * 0.001
}
self.psf_class = PSF(**kwargs_psf)
# 'EXERNAL_SHEAR': external shear
kwargs_shear = {
'gamma1': 0.01,
'gamma2': 0.01
} # gamma_ext: shear strength, psi_ext: shear angel (in radian)
phi, q = 0.2, 0.8
e1, e2 = param_util.phi_q2_ellipticity(phi, q)
kwargs_spemd = {
'theta_E': 1.,
'gamma': 1.8,
'center_x': 0,
'center_y': 0,
'e1': e1,
'e2': e2
}
lens_model_list = ['SPEP', 'SHEAR']
self.kwargs_lens = [kwargs_spemd, kwargs_shear]
self.lens_model_class = LensModel(lens_model_list=lens_model_list)
# list of light profiles (for lens and source)
# 'SERSIC': spherical Sersic profile
kwargs_sersic = {
'amp': 1.,
'R_sersic': 0.1,
'n_sersic': 2,
'center_x': 0,
'center_y': 0
}
# 'SERSIC_ELLIPSE': elliptical Sersic profile
phi, q = 0.2, 0.9
e1, e2 = param_util.phi_q2_ellipticity(phi, q)
kwargs_sersic_ellipse = {
'amp': 1.,
'R_sersic': .6,
'n_sersic': 7,
'center_x': 0,
'center_y': 0,
'e1': e1,
'e2': e2
}
lens_light_model_list = ['SERSIC']
kwargs_lens_light_base = [kwargs_sersic]
lens_light_model_class_base = LightModel(
light_model_list=lens_light_model_list)
source_model_list = ['SERSIC_ELLIPSE']
kwargs_source_base = [kwargs_sersic_ellipse]
source_model_class_base = LightModel(
light_model_list=source_model_list)
self.kwargs_ps = [
{
'ra_source': 0.01,
'dec_source': 0.0,
'source_amp': 1.
}
] # quasar point source position in the source plane and intrinsic brightness
point_source_class_base = PointSource(
point_source_type_list=['SOURCE_POSITION'],
fixed_magnification_list=[True])
kwargs_numerics_base = {
'supersampling_factor': 2,
'supersampling_convolution': False
}
imageModel_base = ImageModel(self.data_class,
self.psf_class,
self.lens_model_class,
source_model_class_base,
lens_light_model_class_base,
point_source_class_base,
kwargs_numerics=kwargs_numerics_base)
image_sim = sim_util.simulate_simple(imageModel_base, self.kwargs_lens,
kwargs_source_base,
kwargs_lens_light_base,
self.kwargs_ps)
self.data_class.update_data(image_sim)
# create a starlet light distributions
n_scales = 6
source_map = imageModel_base.source_surface_brightness(
kwargs_source_base, de_lensed=True, unconvolved=True)
starlets_class = SLIT_Starlets(force_no_pysap=_force_no_pysap)
source_map_starlets = starlets_class.decomposition_2d(
source_map, n_scales)
self.kwargs_source = [{
'amp': source_map_starlets,
'n_scales': n_scales,
'n_pixels': numPix,
'scale': deltaPix,
'center_x': 0,
'center_y': 0
}]
self.source_model_class = LightModel(
light_model_list=['SLIT_STARLETS'])
lens_light_map = imageModel_base.lens_surface_brightness(
kwargs_lens_light_base, unconvolved=True)
starlets_class = SLIT_Starlets(force_no_pysap=_force_no_pysap,
second_gen=True)
lens_light_starlets = starlets_class.decomposition_2d(
lens_light_map, n_scales)
self.kwargs_lens_light = [{
'amp': lens_light_starlets,
'n_scales': n_scales,
'n_pixels': numPix,
'scale': deltaPix,
'center_x': 0,
'center_y': 0
}]
self.lens_light_model_class = LightModel(
light_model_list=['SLIT_STARLETS_GEN2'])
self.kwargs_numerics = {'supersampling_factor': 1}
self.kwargs_pixelbased = {
'supersampling_factor_source':
2, # supersampling of pixelated source grid
# following choices are to minimize pixel solver runtime (not to get accurate reconstruction!)
'threshold_decrease_type': 'none',
'num_iter_source': 2,
'num_iter_lens': 2,
'num_iter_global': 2,
'num_iter_weights': 2,
}
class TestImageModel(object):
"""
tests the source model routines
"""
def setup(self):
self.SimAPI = Simulation()
# data specifics
sigma_bkg = .05 # background noise per pixel
exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
numPix = 100 # cutout pixel size
deltaPix = 0.05 # pixel size in arcsec (area per pixel = deltaPix**2)
fwhm = 0.5 # full width half max of PSF
# PSF specification
kwargs_data = self.SimAPI.data_configure(numPix, deltaPix, exp_time, sigma_bkg)
data_class = Data(kwargs_data)
kwargs_psf = self.SimAPI.psf_configure(psf_type='GAUSSIAN', fwhm=fwhm, kernelsize=31, deltaPix=deltaPix, truncate=3,
kernel=None)
psf_class = PSF(kwargs_psf)
psf_class._psf_error_map = np.zeros_like(psf_class.kernel_point_source)
# 'EXERNAL_SHEAR': external shear
kwargs_shear = {'e1': 0.01, 'e2': 0.01} # gamma_ext: shear strength, psi_ext: shear angel (in radian)
phi, q = 0.2, 0.8
e1, e2 = param_util.phi_q2_ellipticity(phi, q)
kwargs_spemd = {'theta_E': 1., 'gamma': 1.8, 'center_x': 0, 'center_y': 0, 'e1': e1, 'e2': e2}
lens_model_list = ['SPEP', 'SHEAR']
self.kwargs_lens = [kwargs_spemd, kwargs_shear]
lens_model_class = LensModel(lens_model_list=lens_model_list)
# list of light profiles (for lens and source)
# 'SERSIC': spherical Sersic profile
kwargs_sersic = {'amp': 1., 'R_sersic': 0.1, 'n_sersic': 2, 'center_x': 0, 'center_y': 0}
# 'SERSIC_ELLIPSE': elliptical Sersic profile
phi, q = 0.2, 0.9
e1, e2 = param_util.phi_q2_ellipticity(phi, q)
kwargs_sersic_ellipse = {'amp': 1., 'R_sersic': .6, 'n_sersic': 7, 'center_x': 0, 'center_y': 0,
'e1': e1, 'e2': e2}
lens_light_model_list = ['SERSIC']
self.kwargs_lens_light = [kwargs_sersic]
lens_light_model_class = LightModel(light_model_list=lens_light_model_list)
source_model_list = ['SERSIC_ELLIPSE']
self.kwargs_source = [kwargs_sersic_ellipse]
source_model_class = LightModel(light_model_list=source_model_list)
self.kwargs_ps = [{'ra_source': 0.01, 'dec_source': 0.0,
'source_amp': 1.}] # quasar point source position in the source plane and intrinsic brightness
point_source_class = PointSource(point_source_type_list=['SOURCE_POSITION'], fixed_magnification_list=[True])
kwargs_numerics = {'subgrid_res': 2, 'psf_subgrid': True}
imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics)
image_sim = self.SimAPI.simulate(imageModel, self.kwargs_lens, self.kwargs_source,
self.kwargs_lens_light, self.kwargs_ps)
data_class.update_data(image_sim)
self.imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics)
self.solver = LensEquationSolver(lensModel=self.imageModel.LensModel)
def test_source_surface_brightness(self):
source_model = self.imageModel.source_surface_brightness(self.kwargs_source, self.kwargs_lens, unconvolved=False, de_lensed=False)
assert len(source_model) == 100
npt.assert_almost_equal(source_model[10, 10], 0.13939841209844345, decimal=4)
source_model = self.imageModel.source_surface_brightness(self.kwargs_source, self.kwargs_lens, unconvolved=True, de_lensed=False)
assert len(source_model) == 100
npt.assert_almost_equal(source_model[10, 10], 0.13536114618182182, decimal=4)
def test_lens_surface_brightness(self):
lens_flux = self.imageModel.lens_surface_brightness(self.kwargs_lens_light, unconvolved=False)
npt.assert_almost_equal(lens_flux[50, 50], 0.54214440654021534, decimal=4)
lens_flux = self.imageModel.lens_surface_brightness(self.kwargs_lens_light, unconvolved=True)
npt.assert_almost_equal(lens_flux[50, 50], 4.7310552067454452, decimal=4)
def test_image_linear_solve(self):
model, error_map, cov_param, param = self.imageModel.image_linear_solve(self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps, inv_bool=False)
chi2_reduced = self.imageModel.reduced_chi2(model, error_map)
npt.assert_almost_equal(chi2_reduced, 1, decimal=1)
def test_image_with_params(self):
model = self.imageModel.image(self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps, unconvolved=False, source_add=True, lens_light_add=True, point_source_add=True)
error_map = self.imageModel.error_map(self.kwargs_lens, self.kwargs_ps)
chi2_reduced = self.imageModel.reduced_chi2(model, error_map)
npt.assert_almost_equal(chi2_reduced, 1, decimal=1)
def test_point_sources_list(self):
point_source_list = self.imageModel.point_sources_list(self.kwargs_ps, self.kwargs_lens)
assert len(point_source_list) == 4
def test_image_positions(self):
x_im, y_im = self.imageModel.image_positions(self.kwargs_ps, self.kwargs_lens)
ra_pos, dec_pos = self.solver.image_position_from_source(sourcePos_x=self.kwargs_ps[0]['ra_source'],
sourcePos_y=self.kwargs_ps[0]['dec_source'],
kwargs_lens=self.kwargs_lens)
ra_pos_new = x_im[0]
print(ra_pos_new, ra_pos)
npt.assert_almost_equal(ra_pos_new[0], ra_pos[0], decimal=8)
npt.assert_almost_equal(ra_pos_new[1], ra_pos[1], decimal=8)
npt.assert_almost_equal(ra_pos_new[2], ra_pos[2], decimal=8)
npt.assert_almost_equal(ra_pos_new[3], ra_pos[3], decimal=8)
def test_likelihood_data_given_model(self):
logL = self.imageModel.likelihood_data_given_model(self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps, source_marg=False)
npt.assert_almost_equal(logL, -5000, decimal=-3)
logLmarg = self.imageModel.likelihood_data_given_model(self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light,
self.kwargs_ps, source_marg=True)
npt.assert_almost_equal(logL - logLmarg, 0, decimal=-3)
def test_reduced_residuals(self):
model = self.SimAPI.simulate(self.imageModel, self.kwargs_lens, self.kwargs_source,
self.kwargs_lens_light, self.kwargs_ps, no_noise=True)
residuals = self.imageModel.reduced_residuals(model, error_map=0)
npt.assert_almost_equal(np.std(residuals), 1.01, decimal=1)
chi2 = self.imageModel.reduced_chi2(model, error_map=0)
npt.assert_almost_equal(chi2, 1, decimal=1)
def test_numData_evaluate(self):
numData = self.imageModel.numData_evaluate()
assert numData == 10000
def test_fermat_potential(self):
phi_fermat = self.imageModel.fermat_potential(self.kwargs_lens, self.kwargs_ps)
print(phi_fermat)
npt.assert_almost_equal(phi_fermat[0][0], -0.2630531731871062, decimal=3)
npt.assert_almost_equal(phi_fermat[0][1], -0.2809100018126987, decimal=3)
npt.assert_almost_equal(phi_fermat[0][2], -0.5086643370512096, decimal=3)
npt.assert_almost_equal(phi_fermat[0][3], -0.5131716608238992, decimal=3)
def test_add_mask(self):
mask = np.array([[0, 1],[1, 0]])
A = np.ones((10, 4))
A_masked = self.imageModel._add_mask(A, mask)
assert A[0, 1] == A_masked[0, 1]
assert A_masked[0, 3] == 0
def test_point_source_rendering(self):
# initialize data
from lenstronomy.SimulationAPI.simulations import Simulation
SimAPI = Simulation()
numPix = 100
deltaPix = 0.05
kwargs_data = SimAPI.data_configure(numPix, deltaPix, exposure_time=1, sigma_bkg=1)
data_class = Data(kwargs_data)
kernel = np.zeros((5, 5))
kernel[2, 2] = 1
kwargs_psf = {'kernel_point_source': kernel, 'kernel_pixel': kernel, 'psf_type': 'PIXEL'}
psf_class = PSF(kwargs_psf)
lens_model_class = LensModel(['SPEP'])
source_model_class = LightModel([])
lens_light_model_class = LightModel([])
kwargs_numerics = {'subgrid_res': 2, 'point_source_subgrid': 1}
point_source_class = PointSource(point_source_type_list=['LENSED_POSITION'], fixed_magnification_list=[False])
makeImage = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class, point_source_class, kwargs_numerics=kwargs_numerics)
# chose point source positions
x_pix = np.array([10, 5, 10, 90])
y_pix = np.array([40, 50, 60, 50])
ra_pos, dec_pos = makeImage.Data.map_pix2coord(x_pix, y_pix)
e1, e2 = param_util.phi_q2_ellipticity(0, 0.8)
kwargs_lens_init = [{'theta_E': 1, 'gamma': 2, 'e1': e1, 'e2': e2, 'center_x': 0, 'center_y': 0}]
kwargs_else = [{'ra_image': ra_pos, 'dec_image': dec_pos, 'point_amp': np.ones_like(ra_pos)}]
model = makeImage.image(kwargs_lens_init, kwargs_source={}, kwargs_lens_light={}, kwargs_ps=kwargs_else)
image = makeImage.ImageNumerics.array2image(model)
for i in range(len(x_pix)):
npt.assert_almost_equal(image[y_pix[i], x_pix[i]], 1, decimal=2)
x_pix = np.array([10.5, 5.5, 10.5, 90.5])
y_pix = np.array([40, 50, 60, 50])
ra_pos, dec_pos = makeImage.Data.map_pix2coord(x_pix, y_pix)
phi, q = 0., 0.8
e1, e2 = param_util.phi_q2_ellipticity(phi, q)
kwargs_lens_init = [{'theta_E': 1, 'gamma': 2, 'e1': e1, 'e2': e2, 'center_x': 0, 'center_y': 0}]
kwargs_else = [{'ra_image': ra_pos, 'dec_image': dec_pos, 'point_amp': np.ones_like(ra_pos)}]
model = makeImage.image(kwargs_lens_init, kwargs_source={}, kwargs_lens_light={}, kwargs_ps=kwargs_else)
image = makeImage.ImageNumerics.array2image(model)
for i in range(len(x_pix)):
print(int(y_pix[i]), int(x_pix[i]+0.5))
npt.assert_almost_equal(image[int(y_pix[i]), int(x_pix[i])], 0.5, decimal=1)
npt.assert_almost_equal(image[int(y_pix[i]), int(x_pix[i]+0.5)], 0.5, decimal=1)
pointSource = PointSource(point_source_type_list=point_source_list)
kwargs_numerics = {'subgrid_res': 1, 'psf_subgrid': False}
#kwargs_numerics['mask'] = QSO_msk
imageModel = ImageModel(data_class,
psf_class,
source_model_class=lightModel,
point_source_class=pointSource,
kwargs_numerics=kwargs_numerics)
#labels_new = ["Quasar flux"] + ["host{0} flux".format(i) for i in range(3)] + ["host{0} Reff".format(i) for i in range(3)]
for i in range(len(samples_mcmc) / 10):
kwargs_lens_out, kwargs_light_source_out, kwargs_light_lens_out, kwargs_ps_out, kwargs_cosmo = param.args2kwargs(
samples_mcmc[i + len(samples_mcmc) / 10 * 9])
image_reconstructed, _, _, _ = imageModel.image_linear_solve(
kwargs_source=kwargs_light_source_out, kwargs_ps=kwargs_ps_out)
image_ps = imageModel.point_source(kwargs_ps_out)
flux_quasar = np.sum(image_ps)
fluxs, reffs = [], []
for j in range(3):
image_j = imageModel.source_surface_brightness(kwargs_light_source_out,
unconvolved=False,
k=j)
fluxs.append(np.sum(image_j))
reffs.append(kwargs_light_source_out[j]['R_sersic'])
mcmc_new_list.append([flux_quasar] + fluxs + reffs)
if i / 1000 > (i - 1) / 1000:
print "finished translate:", i
print 'The new translated plot:'
plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True)
plt.show()
def fit_qso(QSO_im, psf_ave, psf_std=None, source_params=None,ps_param=None, background_rms=0.04, pix_sz = 0.168,
exp_time = 300., fix_n=None, image_plot = True, corner_plot=True, supersampling_factor = 2,
flux_ratio_plot=False, deep_seed = False, fixcenter = False, QSO_msk=None, QSO_std=None,
tag = None, no_MCMC= False, pltshow = 1, return_Chisq = False, dump_result = False, pso_diag=False):
'''
A quick fit for the QSO image with (so far) single sersice + one PSF. The input psf noise is optional.
Parameter
--------
QSO_im: An array of the QSO image.
psf_ave: The psf image.
psf_std: The psf noise, optional.
source_params: The prior for the source. Default is given. If [], means no Sersic light.
background_rms: default as 0.04
exp_time: default at 2400.
deep_seed: if Ture, more mcmc steps will be performed.
tag: The name tag for save the plot
Return
--------
Will output the fitted image (Set image_plot = True), the corner_plot and the flux_ratio_plot.
source_result, ps_result, image_ps, image_host
To do
--------
'''
# data specifics need to set up based on the data situation
background_rms = background_rms # background noise per pixel (Gaussian)
exp_time = exp_time # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
numPix = len(QSO_im) # cutout pixel size
deltaPix = pix_sz
psf_type = 'PIXEL' # 'gaussian', 'pixel', 'NONE'
kernel = psf_ave
kwargs_numerics = {'supersampling_factor': supersampling_factor, 'supersampling_convolution': False}
if source_params is None:
# here are the options for the host galaxy fitting
fixed_source = []
kwargs_source_init = []
kwargs_source_sigma = []
kwargs_lower_source = []
kwargs_upper_source = []
if fix_n == None:
fixed_source.append({}) # we fix the Sersic index to n=1 (exponential)
kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': 2., 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.})
kwargs_source_sigma.append({'n_sersic': 0.5, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1})
kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': 0.3, 'center_x': -10, 'center_y': -10})
kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3., 'n_sersic': 7., 'center_x': 10, 'center_y': 10})
elif fix_n is not None:
fixed_source.append({'n_sersic': fix_n})
kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': fix_n, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.})
kwargs_source_sigma.append({'n_sersic': 0.001, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1})
kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': fix_n, 'center_x': -10, 'center_y': -10})
kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3, 'n_sersic': fix_n, 'center_x': 10, 'center_y': 10})
source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source]
else:
source_params = source_params
if ps_param is None:
center_x = 0.0
center_y = 0.0
point_amp = QSO_im.sum()/2.
fixed_ps = [{}]
kwargs_ps = [{'ra_image': [center_x], 'dec_image': [center_y], 'point_amp': [point_amp]}]
kwargs_ps_init = kwargs_ps
kwargs_ps_sigma = [{'ra_image': [0.05], 'dec_image': [0.05]}]
kwargs_lower_ps = [{'ra_image': [-0.6], 'dec_image': [-0.6]}]
kwargs_upper_ps = [{'ra_image': [0.6], 'dec_image': [0.6]}]
ps_param = [kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps]
else:
ps_param = ps_param
#==============================================================================
#Doing the QSO fitting
#==============================================================================
kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, background_rms, inverse=True)
data_class = ImageData(**kwargs_data)
kwargs_psf = {'psf_type': psf_type, 'kernel_point_source': kernel}
psf_class = PSF(**kwargs_psf)
data_class.update_data(QSO_im)
point_source_list = ['UNLENSED'] * len(ps_param[0])
pointSource = PointSource(point_source_type_list=point_source_list)
if fixcenter == False:
kwargs_constraints = {'num_point_source_list': [1] * len(ps_param[0])
}
elif fixcenter == True:
kwargs_constraints = {'joint_source_with_point_source': [[i, i] for i in range(len(ps_param[0]))],
'num_point_source_list': [1] * len(ps_param[0])
}
if source_params == []: #fitting image as Point source only.
kwargs_params = {'point_source_model': ps_param}
lightModel = None
kwargs_model = {'point_source_model_list': point_source_list }
imageModel = ImageModel(data_class, psf_class, point_source_class=pointSource, kwargs_numerics=kwargs_numerics)
kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty"
'image_likelihood_mask_list': [QSO_msk]
}
elif source_params != []:
kwargs_params = {'source_model': source_params,
'point_source_model': ps_param}
light_model_list = ['SERSIC_ELLIPSE'] * len(source_params[0])
lightModel = LightModel(light_model_list=light_model_list)
kwargs_model = { 'source_light_model_list': light_model_list,
'point_source_model_list': point_source_list
}
imageModel = ImageModel(data_class, psf_class, source_model_class=lightModel,
point_source_class=pointSource, kwargs_numerics=kwargs_numerics)
# numerical options and fitting sequences
kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty"
'source_marg': False, #In likelihood_module.LikelihoodModule -- whether to fully invert the covariance matrix for marginalization
'check_positive_flux': True,
'image_likelihood_mask_list': [QSO_msk]
}
kwargs_data['image_data'] = QSO_im
if QSO_std is not None:
kwargs_data['noise_map'] = QSO_std
if psf_std is not None:
kwargs_psf['psf_error_map'] = psf_std
image_band = [kwargs_data, kwargs_psf, kwargs_numerics]
multi_band_list = [image_band]
kwargs_data_joint = {'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear'} # 'single-band', 'multi-linear', 'joint-linear'
fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params)
if deep_seed == False:
fitting_kwargs_list = [
['PSO', {'sigma_scale': 0.8, 'n_particles': 100, 'n_iterations': 60}],
['MCMC', {'n_burn': 10, 'n_run': 10, 'walkerRatio': 50, 'sigma_scale': .1}]
]
elif deep_seed == True:
fitting_kwargs_list = [
['PSO', {'sigma_scale': 0.8, 'n_particles': 250, 'n_iterations': 250}],
['MCMC', {'n_burn': 100, 'n_run': 200, 'walkerRatio': 10, 'sigma_scale': .1}]
]
if no_MCMC == True:
fitting_kwargs_list = [fitting_kwargs_list[0],
]
start_time = time.time()
chain_list = fitting_seq.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_seq.best_fit()
ps_result = kwargs_result['kwargs_ps']
source_result = kwargs_result['kwargs_source']
if no_MCMC == False:
sampler_type, samples_mcmc, param_mcmc, dist_mcmc = chain_list[1]
end_time = time.time()
print(end_time - start_time, 'total time needed for computation')
print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ')
imageLinearFit = ImageLinearFit(data_class=data_class, psf_class=psf_class,
source_model_class=lightModel,
point_source_class=pointSource,
kwargs_numerics=kwargs_numerics)
image_reconstructed, error_map, _, _ = imageLinearFit.image_linear_solve(kwargs_source=source_result, kwargs_ps=ps_result)
# this is the linear inversion. The kwargs will be updated afterwards
modelPlot = ModelPlot(multi_band_list, kwargs_model, kwargs_result,
arrow_size=0.02, cmap_string="gist_heat", likelihood_mask_list=[QSO_msk])
image_host = [] #!!! The linear_solver before and after LensModelPlot could have different result for very faint sources.
for i in range(len(source_result)):
image_host.append(imageModel.source_surface_brightness(source_result, de_lensed=True,unconvolved=False,k=i))
image_ps = []
for i in range(len(ps_result)):
image_ps.append(imageModel.point_source(ps_result, k = i))
if pso_diag == True:
f, axes = chain_plot.plot_chain_list(chain_list,0)
if pltshow == 0:
plt.close()
else:
plt.show()
# let's plot the output of the PSO minimizer
reduced_Chisq = imageLinearFit.reduced_chi2(image_reconstructed, error_map)
if image_plot:
f, axes = plt.subplots(3, 3, figsize=(16, 16), sharex=False, sharey=False)
modelPlot.data_plot(ax=axes[0,0], text="Data")
modelPlot.model_plot(ax=axes[0,1])
modelPlot.normalized_residual_plot(ax=axes[0,2], v_min=-6, v_max=6)
modelPlot.decomposition_plot(ax=axes[1,0], text='Host galaxy', source_add=True, unconvolved=True)
modelPlot.decomposition_plot(ax=axes[1,1], text='Host galaxy convolved', source_add=True)
modelPlot.decomposition_plot(ax=axes[1,2], text='All components convolved', source_add=True, lens_light_add=True, point_source_add=True)
modelPlot.subtract_from_data_plot(ax=axes[2,0], text='Data - Point Source', point_source_add=True)
modelPlot.subtract_from_data_plot(ax=axes[2,1], text='Data - host galaxy', source_add=True)
modelPlot.subtract_from_data_plot(ax=axes[2,2], text='Data - host galaxy - Point Source', source_add=True, point_source_add=True)
f.tight_layout()
#f.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0., hspace=0.05)
if tag is not None:
f.savefig('{0}_fitted_image.pdf'.format(tag))
if pltshow == 0:
plt.close()
else:
plt.show()
if corner_plot==True and no_MCMC==False:
# here the (non-converged) MCMC chain of the non-linear parameters
if not samples_mcmc == []:
n, num_param = np.shape(samples_mcmc)
plot = corner.corner(samples_mcmc, labels=param_mcmc, show_titles=True)
if tag is not None:
plot.savefig('{0}_para_corner.pdf'.format(tag))
plt.close()
# if pltshow == 0:
# plt.close()
# else:
# plt.show()
if flux_ratio_plot==True and no_MCMC==False:
param = Param(kwargs_model, kwargs_fixed_source=source_params[2], kwargs_fixed_ps=ps_param[2], **kwargs_constraints)
mcmc_new_list = []
if len(ps_param[2]) == 1:
labels_new = ["Quasar flux"] + ["host{0} flux".format(i) for i in range(len(source_params[0]))]
else:
labels_new = ["Quasar{0} flux".format(i) for i in range(len(ps_param[2]))] + ["host{0} flux".format(i) for i in range(len(source_params[0]))]
if len(samples_mcmc) > 10000:
trans_steps = [len(samples_mcmc)-10000, len(samples_mcmc)]
else:
trans_steps = [0, len(samples_mcmc)]
for i in range(trans_steps[0], trans_steps[1]):
kwargs_out = param.args2kwargs(samples_mcmc[i])
kwargs_light_source_out = kwargs_out['kwargs_source']
kwargs_ps_out = kwargs_out['kwargs_ps']
image_reconstructed, _, _, _ = imageLinearFit.image_linear_solve(kwargs_source=kwargs_light_source_out, kwargs_ps=kwargs_ps_out)
flux_quasar = []
if len(ps_param[0]) == 1:
image_ps_j = imageModel.point_source(kwargs_ps_out)
flux_quasar.append(np.sum(image_ps_j))
else:
for j in range(len(ps_param[0])):
image_ps_j = imageModel.point_source(kwargs_ps_out, k=j)
flux_quasar.append(np.sum(image_ps_j))
fluxs = []
for j in range(len(source_params[0])):
image_j = imageModel.source_surface_brightness(kwargs_light_source_out,unconvolved= False, k=j)
fluxs.append(np.sum(image_j))
mcmc_new_list.append(flux_quasar + fluxs )
if int(i/1000) > int((i-1)/1000) :
print(len(samples_mcmc), "MCMC samplers in total, finished translate:", i )
plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True)
if tag is not None:
plot.savefig('{0}_HOSTvsQSO_corner.pdf'.format(tag))
if pltshow == 0:
plt.close()
else:
plt.show()
if QSO_std is None:
noise_map = np.sqrt(data_class.C_D+np.abs(error_map))
else:
noise_map = np.sqrt(QSO_std**2+np.abs(error_map))
if dump_result == True:
if flux_ratio_plot==True and no_MCMC==False:
trans_paras = [mcmc_new_list, labels_new, 'mcmc_new_list, labels_new']
else:
trans_paras = []
picklename= tag + '.pkl'
best_fit = [source_result, image_host, ps_result, image_ps,'source_result, image_host, ps_result, image_ps']
chain_list_result = [chain_list, 'chain_list']
kwargs_fixed_source=source_params[2]
kwargs_fixed_ps=ps_param[2]
classes = data_class, psf_class, lightModel, pointSource
material = multi_band_list, kwargs_model, kwargs_result, QSO_msk, kwargs_fixed_source, kwargs_fixed_ps, kwargs_constraints, kwargs_numerics, classes
pickle.dump([best_fit, chain_list_result, trans_paras, material], open(picklename, 'wb'))
if return_Chisq == False:
return source_result, ps_result, image_ps, image_host, noise_map
elif return_Chisq == True:
return source_result, ps_result, image_ps, image_host, noise_map, reduced_Chisq
def fit_galaxy(galaxy_im, psf_ave, psf_std=None, source_params=None, background_rms=0.04, pix_sz = 0.08,
exp_time = 300., fix_n=None, image_plot = True, corner_plot=True,
deep_seed = False, galaxy_msk=None, galaxy_std=None, flux_corner_plot = False,
tag = None, no_MCMC= False, pltshow = 1, return_Chisq = False, dump_result = False, pso_diag=False):
'''
A quick fit for the QSO image with (so far) single sersice + one PSF. The input psf noise is optional.
Parameter
--------
galaxy_im: An array of the QSO image.
psf_ave: The psf image.
psf_std: The psf noise, optional.
source_params: The prior for the source. Default is given.
background_rms: default as 0.04
exp_time: default at 2400.
deep_seed: if Ture, more mcmc steps will be performed.
tag: The name tag for save the plot
Return
--------
Will output the fitted image (Set image_plot = True), the corner_plot and the flux_ratio_plot.
source_result, ps_result, image_ps, image_host
To do
--------
'''
# data specifics need to set up based on the data situation
background_rms = background_rms # background noise per pixel (Gaussian)
exp_time = exp_time # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
numPix = len(galaxy_im) # cutout pixel size
deltaPix = pix_sz
if psf_ave is not None:
psf_type = 'PIXEL' # 'gaussian', 'pixel', 'NONE'
kernel = psf_ave
# if psf_std is not None:
# kwargs_numerics = {'subgrid_res': 1, 'psf_error_map': True} #Turn on the PSF error map
# else:
kwargs_numerics = {'supersampling_factor': 1, 'supersampling_convolution': False}
if source_params is None:
# here are the options for the host galaxy fitting
fixed_source = []
kwargs_source_init = []
kwargs_source_sigma = []
kwargs_lower_source = []
kwargs_upper_source = []
# Disk component, as modelled by an elliptical Sersic profile
if fix_n == None:
fixed_source.append({}) # we fix the Sersic index to n=1 (exponential)
kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': 2., 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.})
kwargs_source_sigma.append({'n_sersic': 0.5, 'R_sersic': 0.1, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1})
kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.01, 'n_sersic': 0.3, 'center_x': -10, 'center_y': -10})
kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3., 'n_sersic': 7., 'center_x': 10, 'center_y': 10})
elif fix_n is not None:
fixed_source.append({'n_sersic': fix_n})
kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': fix_n, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.})
kwargs_source_sigma.append({'n_sersic': 0.001, 'R_sersic': 0.1, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1})
kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.01, 'n_sersic': fix_n, 'center_x': -10, 'center_y': -10})
kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3, 'n_sersic': fix_n, 'center_x': 10, 'center_y': 10})
source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source]
else:
source_params = source_params
kwargs_params = {'source_model': source_params}
#==============================================================================
#Doing the QSO fitting
#==============================================================================
kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, exp_time, background_rms, inverse=True)
data_class = ImageData(**kwargs_data)
if psf_ave is not None:
kwargs_psf = {'psf_type': psf_type, 'kernel_point_source': kernel}
else:
kwargs_psf = {'psf_type': 'NONE'}
psf_class = PSF(**kwargs_psf)
data_class.update_data(galaxy_im)
light_model_list = ['SERSIC_ELLIPSE'] * len(source_params[0])
lightModel = LightModel(light_model_list=light_model_list)
kwargs_model = { 'source_light_model_list': light_model_list}
# numerical options and fitting sequences
kwargs_constraints = {}
kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty"
'source_marg': False, #In likelihood_module.LikelihoodModule -- whether to fully invert the covariance matrix for marginalization
'check_positive_flux': True,
'image_likelihood_mask_list': [galaxy_msk]
}
kwargs_data['image_data'] = galaxy_im
if galaxy_std is not None:
kwargs_data['noise_map'] = galaxy_std
if psf_std is not None:
kwargs_psf['psf_error_map'] = psf_std
image_band = [kwargs_data, kwargs_psf, kwargs_numerics]
multi_band_list = [image_band]
kwargs_data_joint = {'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear'} # 'single-band', 'multi-linear', 'joint-linear'
fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params)
if deep_seed == False:
fitting_kwargs_list = [
['PSO', {'sigma_scale': 0.8, 'n_particles': 50, 'n_iterations': 50}],
['MCMC', {'n_burn': 10, 'n_run': 10, 'walkerRatio': 50, 'sigma_scale': .1}]
]
elif deep_seed == True:
fitting_kwargs_list = [
['PSO', {'sigma_scale': 0.8, 'n_particles': 100, 'n_iterations': 80}],
['MCMC', {'n_burn': 10, 'n_run': 15, 'walkerRatio': 50, 'sigma_scale': .1}]
]
elif deep_seed == 'very_deep':
fitting_kwargs_list = [
['PSO', {'sigma_scale': 0.8, 'n_particles': 150, 'n_iterations': 150}],
['MCMC', {'n_burn': 10, 'n_run': 20, 'walkerRatio': 50, 'sigma_scale': .1}]
]
if no_MCMC == True:
fitting_kwargs_list = [fitting_kwargs_list[0],
]
start_time = time.time()
chain_list = fitting_seq.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_seq.best_fit()
ps_result = kwargs_result['kwargs_ps']
source_result = kwargs_result['kwargs_source']
if no_MCMC == False:
sampler_type, samples_mcmc, param_mcmc, dist_mcmc = chain_list[1]
# chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list)
# lens_result, source_result, lens_light_result, ps_result, cosmo_temp = fitting_seq.best_fit()
end_time = time.time()
print(end_time - start_time, 'total time needed for computation')
print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ')
# this is the linear inversion. The kwargs will be updated afterwards
imageModel = ImageModel(data_class, psf_class, source_model_class=lightModel,kwargs_numerics=kwargs_numerics)
imageLinearFit = ImageLinearFit(data_class=data_class, psf_class=psf_class,
source_model_class=lightModel,
kwargs_numerics=kwargs_numerics)
image_reconstructed, error_map, _, _ = imageLinearFit.image_linear_solve(kwargs_source=source_result, kwargs_ps=ps_result)
# image_host = [] #!!! The linear_solver before and after could have different result for very faint sources.
# for i in range(len(source_result)):
# image_host_i = imageModel.source_surface_brightness(source_result,de_lensed=True,unconvolved=False, k=i)
# print("image_host_i", source_result[i])
# print("total flux", image_host_i.sum())
# image_host.append(image_host_i)
# let's plot the output of the PSO minimizer
modelPlot = ModelPlot(multi_band_list, kwargs_model, kwargs_result,
arrow_size=0.02, cmap_string="gist_heat", likelihood_mask_list=[galaxy_msk])
if pso_diag == True:
f, axes = chain_plot.plot_chain_list(chain_list,0)
if pltshow == 0:
plt.close()
else:
plt.show()
reduced_Chisq = imageLinearFit.reduced_chi2(image_reconstructed, error_map)
if image_plot:
f, axes = plt.subplots(1, 3, figsize=(16, 16), sharex=False, sharey=False)
modelPlot.data_plot(ax=axes[0])
modelPlot.model_plot(ax=axes[1])
modelPlot.normalized_residual_plot(ax=axes[2], v_min=-6, v_max=6)
f.tight_layout()
#f.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0., hspace=0.05)
if tag is not None:
f.savefig('{0}_fitted_image.pdf'.format(tag))
if pltshow == 0:
plt.close()
else:
plt.show()
image_host = []
for i in range(len(source_result)):
image_host_i = imageModel.source_surface_brightness(source_result,de_lensed=True,unconvolved=False, k=i)
# print("image_host_i", source_result[i])
# print("total flux", image_host_i.sum())
image_host.append(image_host_i)
if corner_plot==True and no_MCMC==False:
# here the (non-converged) MCMC chain of the non-linear parameters
if not samples_mcmc == []:
n, num_param = np.shape(samples_mcmc)
plot = corner.corner(samples_mcmc, labels=param_mcmc, show_titles=True)
if tag is not None:
plot.savefig('{0}_para_corner.pdf'.format(tag))
if pltshow == 0:
plt.close()
else:
plt.show()
if flux_corner_plot ==True and no_MCMC==False:
param = Param(kwargs_model, kwargs_fixed_source=source_params[2], **kwargs_constraints)
mcmc_new_list = []
labels_new = ["host{0} flux".format(i) for i in range(len(source_params[0]))]
for i in range(len(samples_mcmc)):
kwargs_out = param.args2kwargs(samples_mcmc[i])
kwargs_light_source_out = kwargs_out['kwargs_source']
kwargs_ps_out = kwargs_out['kwargs_ps']
image_reconstructed, _, _, _ = imageLinearFit.image_linear_solve(kwargs_source=kwargs_light_source_out, kwargs_ps=kwargs_ps_out)
fluxs = []
for j in range(len(source_params[0])):
image_j = imageModel.source_surface_brightness(kwargs_light_source_out,unconvolved= False, k=j)
fluxs.append(np.sum(image_j))
mcmc_new_list.append( fluxs )
if int(i/1000) > int((i-1)/1000) :
print(len(samples_mcmc), "MCMC samplers in total, finished translate:", i )
plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True)
if tag is not None:
plot.savefig('{0}_HOSTvsQSO_corner.pdf'.format(tag))
if pltshow == 0:
plt.close()
else:
plt.show()
if galaxy_std is None:
noise_map = np.sqrt(data_class.C_D+np.abs(error_map))
else:
noise_map = np.sqrt(galaxy_std**2+np.abs(error_map))
if dump_result == True:
if flux_corner_plot==True and no_MCMC==False:
trans_paras = [source_params[2], mcmc_new_list, labels_new, 'source_params[2], mcmc_new_list, labels_new']
else:
trans_paras = []
picklename= tag + '.pkl'
best_fit = [source_result, image_host, 'source_result, image_host']
# pso_fit = [chain_list, param_list, 'chain_list, param_list']
# mcmc_fit = [samples_mcmc, param_mcmc, dist_mcmc, 'samples_mcmc, param_mcmc, dist_mcmc']
chain_list_result = [chain_list, 'chain_list']
pickle.dump([best_fit, chain_list_result, trans_paras], open(picklename, 'wb'))
if return_Chisq == False:
return source_result, image_host, noise_map
elif return_Chisq == True:
return source_result, image_host, noise_map, reduced_Chisq
def sim_image(self, info_dict):
"""
Simulate an image based on specifications in sim_dict
Args:
info_dict (dict): A single element from the list produced interanlly by input_reader.Organizer.breakup().
Contains all the properties of a single image to generate.
"""
output_image = []
if self.return_planes:
output_source, output_lens, output_point_source, output_noise = [], [], [], []
output_metadata = []
#set the cosmology
cosmology_info = ['H0', 'Om0', 'Tcmb0', 'Neff', 'm_nu', 'Ob0']
cosmo = FlatLambdaCDM(
**dict_select_choose(list(info_dict.values())[0], cosmology_info))
for band, sim_dict in info_dict.items():
# Parse the info dict
params = self.parse_single_band_info_dict(sim_dict,
cosmo,
band=band)
kwargs_single_band = params[0]
kwargs_model = params[1]
kwargs_numerics = params[2]
kwargs_lens_light_list = params[3]
kwargs_source_list = params[4]
kwargs_point_source_list = params[5]
kwargs_lens_model_list = params[6]
output_metadata += params[7]
# Make image
# data properties
kwargs_data = sim_util.data_configure_simple(
sim_dict['numPix'], kwargs_single_band['pixel_scale'],
kwargs_single_band['exposure_time'])
data_class = ImageData(**kwargs_data)
# psf properties
kwargs_psf = {
'psf_type': kwargs_single_band['psf_type'],
'pixel_size': kwargs_single_band['pixel_scale'],
'fwhm': kwargs_single_band['seeing']
}
psf_class = PSF(**kwargs_psf)
# SimAPI instance for conversion to observed quantities
sim = SimAPI(numpix=sim_dict['numPix'],
kwargs_single_band=kwargs_single_band,
kwargs_model=kwargs_model)
kwargs_lens_model_list = sim.physical2lensing_conversion(
kwargs_mass=kwargs_lens_model_list)
kwargs_lens_light_list, kwargs_source_list, _ = sim.magnitude2amplitude(
kwargs_lens_light_mag=kwargs_lens_light_list,
kwargs_source_mag=kwargs_source_list)
# lens model properties
lens_model_class = LensModel(
lens_model_list=kwargs_model['lens_model_list'],
z_lens=kwargs_model['lens_redshift_list'][0],
z_source=kwargs_model['z_source'],
cosmo=cosmo)
# source properties
source_model_class = LightModel(
light_model_list=kwargs_model['source_light_model_list'])
# lens light properties
lens_light_model_class = LightModel(
light_model_list=kwargs_model['lens_light_model_list'])
# solve for PS positions to incorporate time delays
lensEquationSolver = LensEquationSolver(lens_model_class)
kwargs_ps = []
for ps_idx, ps_mag in enumerate(kwargs_point_source_list):
# modify the SimAPI instance to do one point source at a time
temp_kwargs_model = {k: v for k, v in kwargs_model.items()}
temp_kwargs_model['point_source_model_list'] = [
kwargs_model['point_source_model_list'][ps_idx]
]
sim = SimAPI(numpix=sim_dict['numPix'],
kwargs_single_band=kwargs_single_band,
kwargs_model=temp_kwargs_model)
if kwargs_model['point_source_model_list'][
ps_idx] == 'SOURCE_POSITION':
# convert each image to an amplitude
amplitudes = []
for mag in ps_mag['magnitude']:
ps_dict = {k: v for k, v in ps_mag.items()}
ps_dict['magnitude'] = mag
_, _2, ps = sim.magnitude2amplitude(
kwargs_ps_mag=[ps_dict])
amplitudes.append(ps[0]['source_amp'])
x_image, y_image = lensEquationSolver.findBrightImage(
ps[0]['ra_source'],
ps[0]['dec_source'],
kwargs_lens_model_list,
numImages=4, # max number of images
min_distance=kwargs_single_band['pixel_scale'],
search_window=sim_dict['numPix'] *
kwargs_single_band['pixel_scale'])
magnification = lens_model_class.magnification(
x_image, y_image, kwargs=kwargs_lens_model_list)
#amplitudes = np.array(amplitudes) * np.abs(magnification)
amplitudes = np.array(
[a * m for a, m in zip(amplitudes, magnification)])
kwargs_ps.append({
'ra_image': x_image,
'dec_image': y_image,
'point_amp': amplitudes
})
else:
_, _2, ps = sim.magnitude2amplitude(kwargs_ps_mag=[ps_mag])
kwargs_ps.append(ps[0])
# point source properties
point_source_class = PointSource(point_source_type_list=[
x if x != 'SOURCE_POSITION' else 'LENSED_POSITION'
for x in kwargs_model['point_source_model_list']
],
fixed_magnification_list=[False] *
len(kwargs_ps))
# create an image model
image_model = ImageModel(data_class,
psf_class,
lens_model_class,
source_model_class,
lens_light_model_class,
point_source_class,
kwargs_numerics=kwargs_numerics)
# generate image
image_sim = image_model.image(kwargs_lens_model_list,
kwargs_source_list,
kwargs_lens_light_list, kwargs_ps)
poisson = image_util.add_poisson(
image_sim, exp_time=kwargs_single_band['exposure_time'])
sigma_bkg = data_util.bkg_noise(
kwargs_single_band['read_noise'],
kwargs_single_band['exposure_time'],
kwargs_single_band['sky_brightness'],
kwargs_single_band['pixel_scale'],
num_exposures=kwargs_single_band['num_exposures'])
bkg = image_util.add_background(image_sim, sigma_bkd=sigma_bkg)
image = image_sim + bkg + poisson
# Save theta_E (and sigma_v if used)
for ii in range(len(output_metadata)):
output_metadata.append({
'PARAM_NAME':
output_metadata[ii]['PARAM_NAME'].replace(
'sigma_v', 'theta_E'),
'PARAM_VALUE':
kwargs_lens_model_list[output_metadata[ii]
['LENS_MODEL_IDX']]['theta_E'],
'LENS_MODEL_IDX':
output_metadata[ii]['LENS_MODEL_IDX']
})
# Solve lens equation if desired
if self.solve_lens_equation:
#solver = lens_equation_solver.LensEquationSolver(imSim.LensModel)
#x_mins, y_mins = solver.image_position_from_source(sourcePos_x=kwargs_source_list[0]['center_x'],
# sourcePos_y=kwargs_source_list[0]['center_y'],
# kwargs_lens=kwargs_lens_model_list)
x_mins, y_mins = x_image, y_image
num_source_images = len(x_mins)
# Add noise
image_noise = np.zeros(np.shape(image))
for noise_source_num in range(
1, sim_dict['NUMBER_OF_NOISE_SOURCES'] + 1):
image_noise += self._generate_noise(
sim_dict['NOISE_SOURCE_{0}-NAME'.format(noise_source_num)],
np.shape(image),
select_params(
sim_dict,
'NOISE_SOURCE_{0}-'.format(noise_source_num)))
image += image_noise
# Combine with other bands
output_image.append(image)
# Store plane-separated info if requested
if self.return_planes:
output_lens.append(
image_model.lens_surface_brightness(
kwargs_lens_light_list))
output_source.append(
image_model.source_surface_brightness(
kwargs_source_list, kwargs_lens_model_list))
output_point_source.append(
image_model.point_source(kwargs_ps,
kwargs_lens_model_list))
output_noise.append(image_noise)
# Return the desired information in a dictionary
return_dict = {
'output_image': np.array(output_image),
'output_lens_plane': None,
'output_source_plane': None,
'output_point_source_plane': None,
'output_noise_plane': None,
'x_mins': None,
'y_mins': None,
'num_source_images': None,
'additional_metadata': output_metadata
}
if self.return_planes:
return_dict['output_lens_plane'] = np.array(output_lens)
return_dict['output_source_plane'] = np.array(output_source)
return_dict['output_point_source_plane'] = np.array(
output_point_source)
return_dict['output_noise_plane'] = np.array(output_noise)
if self.solve_lens_equation:
return_dict['x_mins'] = x_mins
return_dict['y_mins'] = y_mins
return_dict['num_source_images'] = num_source_images
return return_dict
