def test_dypolychord(self):
fittingSequence = FittingSequence(self.kwargs_data_joint,
self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
self.kwargs_params)
fitting_list = []
kwargs_dypolychord = {
'sampler_type': 'DYPOLYCHORD',
'kwargs_run': {
'ninit': 8,
'nlive_const': 10,
#'seed_increment': 1,
'resume_dyn_run': False,
#'init_step': 10,
},
'polychord_settings': {
'seed': 1,
#'num_repeats': 20
},
'dypolychord_dynamic_goal':
0.8, # 1 for posterior-only, 0 for evidence-only
'remove_output_dir': True,
}
fitting_list.append(['nested_sampling', kwargs_dypolychord])
chain_list = fittingSequence.fit_sequence(fitting_list)
def initialize_sampler(self, walker_ratio, chains_save_path):
"""
Initialize the sampler to be used by run_samples.
Parameters:
walker_ratio (int): The number of walkers per free parameter.
Must be at least 2.
save_path (str): An h5 path specifying where to save the
sampler chains. If a sampler chain is already present in the
path it will be loaded.
"""
if self.image_selected is False:
raise RuntimeError('Select an image before starting your sampler')
# Set up the fitting sequence and fitting kwargs from lenstronomy
self.walker_ratio = walker_ratio
self.chains_save_path = chains_save_path
ls_kwargs_data_joint = {
'multi_band_list': self.ls_multi_band_list,
'multi_band_type': 'multi-linear'
}
ls_kwargs_constraints = {}
self.fitting_seq = FittingSequence(ls_kwargs_data_joint,
self.ls_kwargs_model,
ls_kwargs_constraints,
self.ls_kwargs_likelihood,
self.ls_kwargs_params)
self.sampler_init = True
def build_fitting_seq(self):
from lenstronomy.Workflow.fitting_sequence import FittingSequence
self.fitting_seq = FittingSequence(self.kwargs_data_joint,
self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
self.kwargs_params)
def test_multinest(self):
# Nested sampler tests
# further decrease the parameter space for nested samplers to run faster
fittingSequence = FittingSequence(self.kwargs_data_joint,
self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
self.kwargs_params)
fitting_list = []
kwargs_update = {
'ps_add_fixed': [[0, ['ra_source', 'dec_source'], [0, 0]]],
'lens_light_add_fixed': [[
0, ['n_sersic', 'R_sersic', 'center_x', 'center_y'],
[4, .1, 0, 0]
]],
'source_add_fixed': [[
0, ['R_sersic', 'e1', 'e2', 'center_x', 'center_y'],
[.6, .1, .1, 0, 0]
]],
'lens_add_fixed': [[
0, ['gamma', 'theta_E', 'e1', 'e2', 'center_x', 'center_y'],
[1.8, 1., .1, .1, 0, 0]
], [1, ['gamma1', 'gamma2'], [0.01, 0.01]]],
'change_source_lower_limit': [[0, ['n_sersic'], [2.9]]],
'change_source_upper_limit': [[0, ['n_sersic'], [3.1]]]
}
fitting_list.append(['update_settings', kwargs_update])
kwargs_multinest = {
'sampler_type': 'MULTINEST',
'kwargs_run': {
'n_live_points': 10,
'evidence_tolerance': 0.5,
'sampling_efficiency':
0.8, # 1 for posterior-only, 0 for evidence-only
'importance_nested_sampling': False,
'multimodal': True,
'const_efficiency_mode':
False, # reduce sampling_efficiency to 5% when True
},
'remove_output_dir': True,
}
fitting_list.append(['nested_sampling', kwargs_multinest])
chain_list2 = fittingSequence.fit_sequence(fitting_list)
kwargs_fixed = fittingSequence._updateManager.fixed_kwargs
npt.assert_almost_equal(kwargs_fixed[0][1]['gamma1'], 0.01, decimal=2)
assert fittingSequence._updateManager._lower_kwargs[1][0][
'n_sersic'] == 2.9
assert fittingSequence._updateManager._upper_kwargs[1][0][
'n_sersic'] == 3.1
kwargs_test = {'kwargs_lens': 1}
fittingSequence.update_state(kwargs_test)
kwargs_out = fittingSequence.best_fit(bijective=True)
assert kwargs_out['kwargs_lens'] == 1
def run_mcmc(self, mcmc_numerics):
"""Sample from the joint likelihood
"""
fitting_seq = FittingSequence(self.kwargs_data_joint,
self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
self.kwargs_params)
fitting_kwargs_list = [['MCMC', mcmc_numerics]]
#with script_utils.HiddenPrints():
chain_list_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list)
kwargs_result_mcmc = fitting_seq.best_fit()
return chain_list_mcmc, kwargs_result_mcmc
def test_fitting_sequence(self):
kwargs_init = [
self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light,
self.kwargs_ps
]
lens_sigma = [{
'theta_E_sigma': 0.1,
'gamma_sigma': 0.1,
'ellipse_sigma': 0.1,
'center_x_sigma': 0.1,
'center_y_sigma': 0.1
}, {
'shear_sigma': 0.1
}]
source_sigma = [{
'R_sersic_sigma': 0.05,
'n_sersic_sigma': 0.5,
'center_x_sigma': 0.1,
'center_y_sigma': 0.1,
'ellipse_sigma': 0.1
}]
lens_light_sigma = [{
'R_sersic_sigma': 0.05,
'n_sersic_sigma': 0.5,
'center_x_sigma': 0.1,
'center_y_sigma': 0.1
}]
ps_sigma = [{'pos_sigma': 1, 'point_amp_sigma': 1}]
kwargs_sigma = [lens_sigma, source_sigma, lens_light_sigma, ps_sigma]
kwargs_fixed = [[{}, {}], [{}], [{}], [{}]]
kwargs_params = [
kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_init, kwargs_init
]
image_band = [self.kwargs_data, self.kwargs_psf, self.kwargs_numerics]
multi_band_list = [image_band]
fittingSequence = FittingSequence(multi_band_list, self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
kwargs_params)
lens_temp, source_temp, lens_light_temp, else_temp, chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fittingSequence.fit_sequence(
fitting_kwargs_list=[])
npt.assert_almost_equal(lens_temp[0]['theta_E'],
self.kwargs_lens[0]['theta_E'],
decimal=2)
n_p = 2
n_i = 2
"""
def test_get_galaxy_galaxy_recipe(self):
"""
Test `get_galaxy_galaxy_recipe` method.
:return:
:rtype:
"""
image = np.random.normal(size=(120, 120))
kwargs_data_joint = {
'multi_band_list': [[{
'image_data':
image,
'background_rms':
0.01,
'exposure_time':
np.ones_like(image),
'ra_at_xy_0':
0.,
'dec_at_xy_0':
0.,
'transform_pix2angle':
np.array([[-0.01, 0], [0, 0.01]])
}, {}, {}]],
'multi_band_type':
'multi-linear'
}
fitting_kwargs_list = self.recipe.get_galaxy_galaxy_recipe(
kwargs_data_joint)
assert isinstance(fitting_kwargs_list, list)
# test the recipe by running it fully
config = deepcopy(self.config)
config.settings['model']['source_light'] = ['SHAPELETS']
recipe = Recipe(config)
fitting_sequence = FittingSequence(
kwargs_data_joint,
config.get_kwargs_model(),
config.get_kwargs_constraints(),
config.get_kwargs_likelihood(),
config.get_kwargs_params(),
)
fitting_kwargs_list = recipe.get_recipe(
kwargs_data_joint=kwargs_data_joint, recipe_name='galaxy-galaxy')
fitting_sequence.fit_sequence(fitting_kwargs_list)
def test_fitting_sequence(self):
#kwargs_init = [self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps]
lens_sigma = [{'theta_E': 0.1, 'gamma': 0.1, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1}, {'e1': 0.1, 'e2': 0.1}]
lens_lower = [{'theta_E': 0., 'gamma': 1.5, 'center_x': -2, 'center_y': -2, 'e1': -0.4, 'e2': -0.4}, {'e1': -0.3, 'e2': -0.3}]
lens_upper = [{'theta_E': 10., 'gamma': 2.5, 'center_x': 2, 'center_y': 2, 'e1': 0.4, 'e2': 0.4}, {'e1': 0.3, 'e2': 0.3}]
source_sigma = [{'R_sersic': 0.05, 'n_sersic': 0.5, 'center_x': 0.1, 'center_y': 0.1, 'e1': 0.1, 'e2': 0.1}]
source_lower = [{'R_sersic': 0.01, 'n_sersic': 0.5, 'center_x': -2, 'center_y': -2, 'e1': -0.4, 'e2': -0.4}]
source_upper = [{'R_sersic': 10, 'n_sersic': 5.5, 'center_x': 2, 'center_y': 2, 'e1': 0.4, 'e2': 0.4}]
lens_light_sigma = [{'R_sersic': 0.05, 'n_sersic': 0.5, 'center_x': 0.1, 'center_y': 0.1}]
lens_light_lower = [{'R_sersic': 0.01, 'n_sersic': 0.5, 'center_x': -2, 'center_y': -2}]
lens_light_upper = [{'R_sersic': 10, 'n_sersic': 5.5, 'center_x': 2, 'center_y': 2}]
ps_sigma = [{'ra_source': 1, 'dec_source': 1, 'point_amp': 1}]
lens_param = self.kwargs_lens, lens_sigma, [{}, {}], lens_lower, lens_upper
source_param = self.kwargs_source, source_sigma, [{}], source_lower, source_upper
lens_light_param = self.kwargs_lens_light, lens_light_sigma, [{}], lens_light_lower, lens_light_upper
ps_param = self.kwargs_ps, ps_sigma, [{}], self.kwargs_ps, self.kwargs_ps
kwargs_params = {'lens_model': lens_param,
'source_model': source_param,
'lens_light_model': lens_light_param,
'point_source_model': ps_param,
#'cosmography': cosmo_param
}
#kwargs_params = [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_init, kwargs_init]
image_band = [self.kwargs_data, self.kwargs_psf, self.kwargs_numerics]
multi_band_list = [image_band]
fittingSequence = FittingSequence(multi_band_list, self.kwargs_model, self.kwargs_constraints, self.kwargs_likelihood, kwargs_params)
lens_temp, source_temp, lens_light_temp, else_temp, cosmo_temp, chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fittingSequence.fit_sequence(fitting_kwargs_list=[])
npt.assert_almost_equal(lens_temp[0]['theta_E'], self.kwargs_lens[0]['theta_E'], decimal=2)
n_p = 2
n_i = 2
fitting_kwargs_list = [
{'fitting_routine': 'PSO', 'sigma_scale': 1, 'n_particles': n_p, 'n_iterations': n_i},
{'fitting_routine': 'MCMC', 'sigma_scale': 0.1, 'n_burn': 1, 'n_run': 1, 'walkerRatio': 2},
{'fitting_routine': 'align_images', 'lower_limit_shift': -0.1, 'upper_limit_shift': 0.1, 'n_particles': 2, 'n_iterations': 2},
{'fitting_routine': 'psf_iteration', 'psf_iter_num': 2, 'psf_iter_factor': 0.5, 'kwargs_psf_iter': {'stacking_option': 'mean'}}
]
lens_temp, source_temp, lens_light_temp, else_temp, cosmo_temp, chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fittingSequence.fit_sequence(fitting_kwargs_list=fitting_kwargs_list)
npt.assert_almost_equal(lens_temp[0]['theta_E'], self.kwargs_lens[0]['theta_E'], decimal=1)
def test_nautilus(self):
kwargs_params = copy.deepcopy(self.kwargs_params)
fittingSequence = FittingSequence(self.kwargs_data_joint,
self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
kwargs_params)
fitting_list = []
kwargs_nautilus = {
'prior_type': 'uniform',
'thread_count': 1,
'verbose': True,
'one_step': True,
'n_live': 2,
'random_state': 42
}
fitting_list.append(['Nautilus', kwargs_nautilus])
chain_list = fittingSequence.fit_sequence(fitting_list)
def test_dynesty(self):
kwargs_params = copy.deepcopy(self.kwargs_params)
kwargs_params['lens_model'][0][0]['theta_E'] += 0.01
fittingSequence = FittingSequence(self.kwargs_data_joint,
self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
kwargs_params)
fitting_list = []
kwargs_dynesty = {
'sampler_type': 'DYNESTY',
'kwargs_run': {
'dlogz_init': 0.01,
'nlive_init': 6,
'nlive_batch': 6,
'maxbatch': 1,
},
}
fitting_list.append(['nested_sampling', kwargs_dynesty])
chain_list = fittingSequence.fit_sequence(fitting_list)
def __init__(self, kwargs_data_joint, kwargs_model,lens_params,source_params,
lenslight_params=None, kwargs_constraints=None, kwargs_likelihood=None):
"""
class to manage cluster source reconstruction.
This class inherited the FittingSequence class in Workflow module of lenstronomy.
:param kwargs_data_joint: keywords arguments of [data, psf, numericals] in lenstronomy convention.
:param kwargs_model: name of model list
:param lens_params: lens model keywords arguments [kwargs_lens_init, kwargs_lens_sigma, kwargs_fixed_lens, kwargs_lower_lens, kwargs_upper_lens]
:param source_params: source model keywords arguments [kwargs_source_init, kwargs_source_sigma, kwargs_fixed_source, kwargs_lower_source, kwargs_upper_source]
:param kwargs_constraints: contraints on models
:param kwargs_likelihood: options of calculating likelihood, see more: LikelihoodModule class in Sampling module of lenstronomy.
"""
self.kwargs_data_joint =kwargs_data_joint
self.multi_band_list = kwargs_data_joint.get('multi_band_list', [])
self.kwargs_model =kwargs_model
kwargs_params = {'lens_model': lens_params, 'source_model': source_params, 'lens_light_model': lenslight_params}
self.kwargs_params= kwargs_params
if kwargs_constraints is None:
kwargs_constraints ={}
if kwargs_likelihood is None:
kwargs_likelihood = {'source_marg': False, 'check_positive_flux': True}
self.fitting_seq_src = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params)
def test_zeus(self):
# we make a very basic lens+source model to feed to check zeus can be run through fitting sequence
# we don't use the kwargs defined in setup() as those are modified during the tests; using unique kwargs here is safer
# data specifics
sigma_bkg = 0.05 # background noise per pixel
exp_time = 100 # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
numPix = 10 # 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)
data_class = ImageData(**kwargs_data)
kwargs_psf_gaussian = {
'psf_type': 'GAUSSIAN',
'fwhm': fwhm,
'pixel_size': deltaPix,
'truncation': 3
}
psf_gaussian = PSF(**kwargs_psf_gaussian)
kwargs_psf = {
'psf_type': 'PIXEL',
'kernel_point_source': psf_gaussian.kernel_point_source,
'psf_error_map': np.zeros_like(psf_gaussian.kernel_point_source)
}
psf_class = PSF(**kwargs_psf)
# make a lens
lens_model_list = ['EPL']
kwargs_epl = {
'theta_E': 0.6,
'gamma': 2.6,
'center_x': 0.0,
'center_y': 0.0,
'e1': 0.1,
'e2': 0.1
}
kwargs_lens = [kwargs_epl]
lens_model_class = LensModel(lens_model_list=lens_model_list)
# make a source
source_model_list = ['SERSIC_ELLIPSE']
kwargs_sersic_ellipse = {
'amp': 1.,
'R_sersic': 0.6,
'n_sersic': 3,
'center_x': 0.0,
'center_y': 0.0,
'e1': 0.1,
'e2': 0.1
}
kwargs_source = [kwargs_sersic_ellipse]
source_model_class = LightModel(light_model_list=source_model_list)
kwargs_numerics = {
'supersampling_factor': 1,
'supersampling_convolution': False
}
imageModel = ImageModel(data_class,
psf_class,
lens_model_class,
source_model_class,
kwargs_numerics=kwargs_numerics)
image_sim = sim_util.simulate_simple(imageModel, kwargs_lens,
kwargs_source)
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
}
lens_fixed = [{}]
lens_sigma = [{
'theta_E': 0.1,
'gamma': 0.1,
'e1': 0.1,
'e2': 0.1,
'center_x': 0.1,
'center_y': 0.1
}]
lens_lower = [{
'theta_E': 0.,
'gamma': 1.5,
'center_x': -2,
'center_y': -2,
'e1': -0.4,
'e2': -0.4
}]
lens_upper = [{
'theta_E': 10.,
'gamma': 2.5,
'center_x': 2,
'center_y': 2,
'e1': 0.4,
'e2': 0.4
}]
source_fixed = [{}]
source_sigma = [{
'R_sersic': 0.05,
'n_sersic': 0.5,
'center_x': 0.1,
'center_y': 0.1,
'e1': 0.1,
'e2': 0.1
}]
source_lower = [{
'R_sersic': 0.01,
'n_sersic': 0.5,
'center_x': -2,
'center_y': -2,
'e1': -0.4,
'e2': -0.4
}]
source_upper = [{
'R_sersic': 10,
'n_sersic': 5.5,
'center_x': 2,
'center_y': 2,
'e1': 0.4,
'e2': 0.4
}]
lens_param = [
kwargs_lens, lens_sigma, lens_fixed, lens_lower, lens_upper
]
source_param = [
kwargs_source, source_sigma, source_fixed, source_lower,
source_upper
]
kwargs_params = {
'lens_model': lens_param,
'source_model': source_param
}
kwargs_constraints = {}
multi_band_list = [[kwargs_data, kwargs_psf, kwargs_numerics]]
kwargs_data_joint = {
'multi_band_list': multi_band_list,
'multi_band_type': 'multi-linear'
}
kwargs_likelihood = {'source_marg': True}
fittingSequence = FittingSequence(kwargs_data_joint, kwargs_model,
kwargs_constraints,
kwargs_likelihood, kwargs_params)
fitting_list = []
kwargs_zeus = {
'sampler_type': 'ZEUS',
'n_burn': 2,
'n_run': 2,
'walkerRatio': 4
}
fitting_list.append(['MCMC', kwargs_zeus])
chain_list = fittingSequence.fit_sequence(fitting_list)
def test_fitting_sequence(self):
fittingSequence = FittingSequence(self.kwargs_data_joint,
self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
self.kwargs_params)
kwargs_result = fittingSequence.best_fit(bijective=False)
lens_temp = kwargs_result['kwargs_lens']
npt.assert_almost_equal(lens_temp[0]['theta_E'],
self.kwargs_lens[0]['theta_E'],
decimal=2)
logL = fittingSequence.best_fit_likelihood
print(logL, 'test')
#print(lens_temp, source_temp, lens_light_temp, ps_temp, cosmo_temp)
assert logL < 0
bic = fittingSequence.bic
assert bic > 0
#npt.assert_almost_equal(bic, 20000000220.29376, decimal=-4)
#npt.assert_almost_equal(logL, -10000000061.792593, decimal=-4)
n_p = 2
n_i = 2
fitting_list = []
kwargs_pso = {
'sigma_scale': 1,
'n_particles': n_p,
'n_iterations': n_i
}
fitting_list.append(['PSO', kwargs_pso])
kwargs_align = {
'lowerLimit': -0.1,
'upperLimit': 0.1,
'n_particles': 2,
'n_iterations': 2
}
fitting_list.append(['align_images', kwargs_align])
kwargs_psf_iter = {
'num_iter': 2,
'psf_iter_factor': 0.5,
'stacking_method': 'mean',
'new_procedure': False
}
fitting_list.append(['psf_iteration', kwargs_psf_iter])
fitting_list.append(['restart', None])
fitting_list.append(['fix_not_computed', {'free_bands': [True]}])
n_sersic_overwrite = 4
kwargs_update = {
'lens_light_add_fixed': [[0, ['n_sersic'], [n_sersic_overwrite]]],
'lens_light_remove_fixed': [[0, ['center_x']]],
'change_source_lower_limit': [[0, ['n_sersic'], [0.1]]],
'change_source_upper_limit': [[0, ['n_sersic'], [10]]]
}
fitting_list.append(['update_settings', kwargs_update])
chain_list = fittingSequence.fit_sequence(fitting_list)
lens_fixed, source_fixed, lens_light_fixed, ps_fixed, special_fixed, extinction_fixed = fittingSequence._updateManager.fixed_kwargs
kwargs_result = fittingSequence.best_fit(bijective=False)
npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'],
self.kwargs_lens[0]['theta_E'],
decimal=1)
npt.assert_almost_equal(
fittingSequence._updateManager._lens_light_fixed[0]['n_sersic'],
n_sersic_overwrite,
decimal=8)
npt.assert_almost_equal(lens_light_fixed[0]['n_sersic'], 4, decimal=-1)
assert fittingSequence._updateManager._lower_kwargs[1][0][
'n_sersic'] == 0.1
assert fittingSequence._updateManager._upper_kwargs[1][0][
'n_sersic'] == 10
# test 'set_param_value' fitting sequence
fitting_list = [[
'set_param_value', {
'lens': [[1, ['gamma1'], [0.013]]]
}
], ['set_param_value', {
'lens_light': [[0, ['center_x'], [0.009]]]
}], ['set_param_value', {
'source': [[0, ['n_sersic'], [2.993]]]
}], ['set_param_value', {
'ps': [[0, ['ra_source'], [0.007]]]
}]]
fittingSequence.fit_sequence(fitting_list)
kwargs_set = fittingSequence._updateManager.parameter_state
assert kwargs_set['kwargs_lens'][1]['gamma1'] == 0.013
assert kwargs_set['kwargs_lens_light'][0]['center_x'] == 0.009
assert kwargs_set['kwargs_source'][0]['n_sersic'] == 2.993
assert kwargs_set['kwargs_ps'][0]['ra_source'] == 0.007
def make_lensmodel(lens_info, theta_E, source_info, box_f):
# lens data specifics
lens_image = lens_info['image']
psf_lens = lens_info['psf']
background_rms = background_rms_image(5, lens_image)
exposure_time = 100
kwargs_data_lens = sim_util.data_configure_simple(len(lens_image),
lens_info['deltapix'],
exposure_time,
background_rms)
kwargs_data_lens['image_data'] = lens_image
data_class_lens = ImageData(**kwargs_data_lens)
#PSF
kwargs_psf_lens = {
'psf_type': 'PIXEL',
'pixel_size': lens_info['deltapix'],
'kernel_point_source': psf_lens
}
psf_class_lens = PSF(**kwargs_psf_lens)
# lens light model
lens_light_model_list = ['SERSIC_ELLIPSE']
lens_light_model_class = LightModel(light_model_list=lens_light_model_list)
kwargs_model = {'lens_light_model_list': lens_light_model_list}
kwargs_numerics_galfit = {'supersampling_factor': 1}
kwargs_constraints = {}
kwargs_likelihood = {'check_bounds': True}
image_band = [kwargs_data_lens, kwargs_psf_lens, kwargs_numerics_galfit]
multi_band_list = [image_band]
kwargs_data_joint = {
'multi_band_list': multi_band_list,
'multi_band_type': 'multi-linear'
}
# Sersic component
fixed_lens_light = [{}]
kwargs_lens_light_init = [{
'R_sersic': .1,
'n_sersic': 4,
'e1': 0,
'e2': 0,
'center_x': 0,
'center_y': 0
}]
kwargs_lens_light_sigma = [{
'n_sersic': 0.5,
'R_sersic': 0.2,
'e1': 0.1,
'e2': 0.1,
'center_x': 0.1,
'center_y': 0.1
}]
kwargs_lower_lens_light = [{
'e1': -0.5,
'e2': -0.5,
'R_sersic': 0.01,
'n_sersic': 0.5,
'center_x': -10,
'center_y': -10
}]
kwargs_upper_lens_light = [{
'e1': 0.5,
'e2': 0.5,
'R_sersic': 10,
'n_sersic': 8,
'center_x': 10,
'center_y': 10
}]
lens_light_params = [
kwargs_lens_light_init, kwargs_lens_light_sigma, fixed_lens_light,
kwargs_lower_lens_light, kwargs_upper_lens_light
]
kwargs_params = {'lens_light_model': lens_light_params}
fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model,
kwargs_constraints, kwargs_likelihood,
kwargs_params)
fitting_kwargs_list = [[
'PSO', {
'sigma_scale': 1.,
'n_particles': 50,
'n_iterations': 50
}
]]
chain_list = fitting_seq.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_seq.best_fit()
modelPlot = ModelPlot(multi_band_list, kwargs_model, kwargs_result)
# Lens light best result
kwargs_light_lens = kwargs_result['kwargs_lens_light'][0]
#Lens model
kwargs_lens_list = [{
'theta_E': theta_E,
'e1': kwargs_light_lens['e1'],
'e2': kwargs_light_lens['e2'],
'center_x': kwargs_light_lens['center_x'],
'center_y': kwargs_light_lens['center_y']
}]
lensModel = LensModel(['SIE'])
lme = LensModelExtensions(lensModel)
#random position for the source
x_crit_list, y_crit_list = lme.critical_curve_tiling(
kwargs_lens_list,
compute_window=(len(source_info['image'])) * (source_info['deltapix']),
start_scale=source_info['deltapix'],
max_order=10)
if len(x_crit_list) > 2 and len(y_crit_list) > 2:
x_caustic_list, y_caustic_list = lensModel.ray_shooting(
x_crit_list, y_crit_list, kwargs_lens_list)
xsamp0 = np.arange(
min(x_caustic_list) - min(x_caustic_list) * box_f[0],
max(x_caustic_list) + max(x_caustic_list) * box_f[1], 0.1)
xsamp = xsamp0[abs(xsamp0.round(1)) != 0.1]
ysamp0 = np.arange(
min(y_caustic_list) - min(y_caustic_list) * box_f[0],
max(y_caustic_list) + max(y_caustic_list) * box_f[1], 0.1)
ysamp = ysamp0[abs(ysamp0.round(1)) != 0.1]
if len(xsamp) == 0 or len(ysamp) == 0:
x_shift, y_shift = 0.15, 0.15 #arcseconds
else:
y_shift = rand.sample(list(ysamp), 1)[0]
x_shift = rand.sample(list(xsamp), 1)[0]
else:
x_shift, y_shift = -0.15, 0.15 #arcseconds
x_caustic_list = [0]
y_caustic_list = [0]
solver = LensEquationSolver(lensModel)
theta_ra, theta_dec = solver.image_position_from_source(
x_shift, y_shift, kwargs_lens_list)
if len(theta_ra) <= 1:
x_shift, y_shift = -0.2, -0.2 #arcseconds1
if abs(x_shift) >= int(theta_E) or abs(y_shift) >= int(theta_E):
x_shift, y_shift = 0.3, -0.3
print('BLABLA')
print('HERE',
min(x_caustic_list) - min(x_caustic_list) * box_f[0],
max(x_caustic_list) + max(x_caustic_list) * box_f[1],
min(y_caustic_list) - min(y_caustic_list) * box_f[0],
max(y_caustic_list) + max(y_caustic_list) * box_f[1])
return {
'lens_light_model_list': ['SERSIC_ELLIPSE'],
'kwargs_light_lens': [kwargs_light_lens],
'lens_light_model_class': lens_light_model_class,
'kwargs_lens_list': kwargs_lens_list,
'kwargs_data_lens': kwargs_data_lens,
'source_shift': [x_shift, y_shift]
}
def main():
args = parse_args()
test_cfg = TestConfig.from_file(args.test_config_file_path)
train_val_cfg = TrainValConfig.from_file(test_cfg.train_val_config_file_path)
baobab_cfg = get_baobab_config(test_cfg.data.test_dir)
# Set device and default data type
device = torch.device(test_cfg.device_type)
if device.type == 'cuda':
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
seed_everything(test_cfg.global_seed)
############
# Data I/O #
############
test_data = XYCosmoData(test_cfg.data.test_dir, data_cfg=train_val_cfg.data)
master_truth = test_data.cosmo_df
master_truth = metadata_utils.add_qphi_columns(master_truth)
master_truth = metadata_utils.add_gamma_psi_ext_columns(master_truth)
if test_cfg.data.lens_indices is None:
if args.lens_indices_path is None:
# Test on all n_test lenses in the test set
n_test = test_cfg.data.n_test
lens_range = range(n_test)
else:
# Test on the lens indices in a text file at the specified path
lens_range = []
with open(args.lens_indices_path, "r") as f:
for line in f:
lens_range.append(int(line.strip()))
n_test = len(lens_range)
print("Performing H0 inference on {:d} specified lenses...".format(n_test))
else:
if args.lens_indices_path is None:
# Test on the lens indices specified in the test config file
lens_range = test_cfg.data.lens_indices
n_test = len(lens_range)
print("Performing H0 inference on {:d} specified lenses...".format(n_test))
else:
raise ValueError("Specific lens indices were specified in both the test config file and the command-line argument.")
batch_size = max(lens_range) + 1
# Output directory into which the H0 histograms and H0 samples will be saved
out_dir = test_cfg.out_dir
if not os.path.exists(out_dir):
os.makedirs(out_dir)
print("Destination folder path: {:s}".format(out_dir))
else:
raise OSError("Destination folder already exists.")
######################
# Load trained state #
######################
# Instantiate loss function, to append to the MCMC objective as the prior
orig_Y_cols = train_val_cfg.data.Y_cols
# Instantiate MCMC parameter penalty function
params_to_remove = ['lens_light_R_sersic'] #'src_light_R_sersic']
mcmc_Y_cols = [col for col in orig_Y_cols if col not in params_to_remove]
mcmc_Y_dim = len(mcmc_Y_cols)
null_spread = True
#init_D_dt = np.random.uniform(0.0, 10000.0, size=(batch_size, n_walkers, 1)) # FIXME: init H0 hardcoded
kwargs_model = dict(lens_model_list=['PEMD', 'SHEAR'],
point_source_model_list=['SOURCE_POSITION'],
source_light_model_list=['SERSIC_ELLIPSE'])
astro_sig = test_cfg.image_position_likelihood.sigma
# Get H0 samples for each system
if not test_cfg.time_delay_likelihood.baobab_time_delays:
if 'abcd_ordering_i' not in master_truth:
raise ValueError("If the time delay measurements were not generated using Baobab, the user must specify the order of image positions in which the time delays are listed, in order of increasing dec.")
kwargs_lens_eq_solver = {'min_distance': 0.05, 'search_window': baobab_cfg.instrument.pixel_scale*baobab_cfg.image.num_pix, 'num_iter_max': 100}
#n_walkers = test_cfg.numerics.mcmc.walkerRatio*(mcmc_Y_dim + 1) # BNN params + H0 times walker ratio
#init_pos = np.tile(master_truth[mcmc_Y_cols].iloc[:batch_size].values[:, np.newaxis, :], [1, n_walkers, 1])
#init_D_dt = np.random.uniform(0.0, 10000.0, size=(batch_size, n_walkers, 1))
#print(init_pos.shape, init_D_dt.shape)
total_progress = tqdm(total=n_test)
# For each lens system...
for i, lens_i in enumerate(lens_range):
# Each lens gets a unique random state for td and vd measurement error realizations.
rs_lens = np.random.RandomState(lens_i)
###########################
# Relevant data and prior #
###########################
data_i = master_truth.iloc[lens_i].copy()
# Init values for the lens model params
init_info = dict(zip(mcmc_Y_cols, data_i[mcmc_Y_cols].values)) # truth params
lcdm = LCDM(z_lens=data_i['z_lens'], z_source=data_i['z_src'], flat=True)
true_img_dec = np.array(literal_eval(data_i['y_image']))
n_img = len(true_img_dec)
true_td = np.array(literal_eval(data_i['true_td']))
measured_td = true_td + rs_lens.randn(*true_td.shape)*test_cfg.error_model.time_delay_error
measured_td_sig = test_cfg.time_delay_likelihood.sigma # np.ones(n_img - 1)*
measured_img_dec = true_img_dec + rs_lens.randn(n_img)*astro_sig
increasing_dec_i = np.argsort(true_img_dec) #np.argsort(measured_img_dec)
measured_td = h0_utils.reorder_to_tdlmc(measured_td, increasing_dec_i, range(n_img)) # need to use measured dec to order
measured_img_dec = h0_utils.reorder_to_tdlmc(measured_img_dec, increasing_dec_i, range(n_img))
measured_td_wrt0 = measured_td[1:] - measured_td[0]
kwargs_data_joint = dict(time_delays_measured=measured_td_wrt0,
time_delays_uncertainties=measured_td_sig,
)
#############################
# Parameter init and bounds #
#############################
lens_kwargs = mcmc_utils.get_lens_kwargs(init_info, null_spread=null_spread)
ps_kwargs = mcmc_utils.get_ps_kwargs_src_plane(init_info, astro_sig, null_spread=null_spread)
src_light_kwargs = mcmc_utils.get_light_kwargs(init_info['src_light_R_sersic'], null_spread=null_spread)
special_kwargs = mcmc_utils.get_special_kwargs(n_img, astro_sig, D_dt_sigma=2000, null_spread=null_spread) # image position offset and time delay distance, aka the "special" parameters
kwargs_params = {'lens_model': lens_kwargs,
'point_source_model': ps_kwargs,
'source_model': src_light_kwargs,
'special': special_kwargs,}
if test_cfg.numerics.solver_type == 'NONE':
solver_type = 'NONE'
else:
solver_type = 'PROFILE_SHEAR' if n_img == 4 else 'CENTER'
#solver_type = 'NONE'
kwargs_constraints = {'num_point_source_list': [n_img],
'Ddt_sampling': True,
'solver_type': solver_type,}
kwargs_likelihood = {'time_delay_likelihood': True,
'sort_images_by_dec': True,
'prior_lens': [],
'prior_special': [],
'check_bounds': True,
'check_matched_source_position': False,
'source_position_tolerance': 0.01,
'source_position_sigma': 0.01,
'source_position_likelihood': False,
'custom_logL_addition': None,
'kwargs_lens_eq_solver': kwargs_lens_eq_solver}
###########################
# MCMC posterior sampling #
###########################
fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params, verbose=False, mpi=False)
if i == 0:
param_class = fitting_seq._updateManager.param_class
n_params, param_class_Y_cols = param_class.num_param()
#init_pos = mcmc_utils.reorder_to_param_class(mcmc_Y_cols, param_class_Y_cols, init_pos, init_D_dt)
# MCMC sample from the post-processed BNN posterior jointly with cosmology
lens_i_start_time = time.time()
#test_cfg.numerics.mcmc.update(init_samples=init_pos[lens_i, :, :])
fitting_kwargs_list_mcmc = [['MCMC', test_cfg.numerics.mcmc]]
#with HiddenPrints():
#try:
chain_list_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list_mcmc)
kwargs_result_mcmc = fitting_seq.best_fit()
#except:
# print("lens {:d} skipped".format(lens_i))
# total_progress.update(1)
# continue
lens_i_end_time = time.time()
inference_time = (lens_i_end_time - lens_i_start_time)/60.0 # min
#############################
# Plotting the MCMC samples #
#############################
# sampler_type : 'EMCEE'
# samples_mcmc : np.array of shape `[n_mcmc_eval, n_params]`
# param_mcmc : list of str of length n_params, the parameter names
sampler_type, samples_mcmc, param_mcmc, _ = chain_list_mcmc[0]
new_samples_mcmc = mcmc_utils.postprocess_mcmc_chain(kwargs_result_mcmc, samples_mcmc, kwargs_model, lens_kwargs[2], ps_kwargs[2], src_light_kwargs[2], special_kwargs[2], kwargs_constraints)
# Plot D_dt histogram
D_dt_samples = new_samples_mcmc['D_dt'].values
true_D_dt = lcdm.D_dt(H_0=data_i['H0'], Om0=0.3)
data_i['D_dt'] = true_D_dt
# Export D_dt samples for this lens
lens_inference_dict = dict(
D_dt_samples=D_dt_samples, # kappa_ext=0 for these samples
inference_time=inference_time,
true_D_dt=true_D_dt,
)
lens_inference_dict_save_path = os.path.join(out_dir, 'D_dt_dict_{0:04d}.npy'.format(lens_i))
np.save(lens_inference_dict_save_path, lens_inference_dict)
# Optionally export the MCMC samples
if test_cfg.export.mcmc_samples:
mcmc_samples_path = os.path.join(out_dir, 'mcmc_samples_{0:04d}.csv'.format(lens_i))
new_samples_mcmc.to_csv(mcmc_samples_path, index=None)
# Optionally export the D_dt histogram
if test_cfg.export.D_dt_histogram:
cleaned_D_dt_samples = h0_utils.remove_outliers_from_lognormal(D_dt_samples, 3)
_ = plotting_utils.plot_D_dt_histogram(cleaned_D_dt_samples, lens_i, true_D_dt, save_dir=out_dir)
# Optionally export the plot of MCMC chain
if test_cfg.export.mcmc_chain:
mcmc_chain_path = os.path.join(out_dir, 'mcmc_chain_{0:04d}.png'.format(lens_i))
plotting_utils.plot_mcmc_chain(chain_list_mcmc, mcmc_chain_path)
# Optionally export posterior cornerplot of select lens model parameters with D_dt
if test_cfg.export.mcmc_corner:
mcmc_corner_path = os.path.join(out_dir, 'mcmc_corner_{0:04d}.png'.format(lens_i))
plotting_utils.plot_mcmc_corner(new_samples_mcmc[test_cfg.export.mcmc_cols], data_i[test_cfg.export.mcmc_cols], test_cfg.export.mcmc_col_labels, mcmc_corner_path)
total_progress.update(1)
gc.collect()
total_progress.close()
kwargs_lens_light_init = []
kwargs_lens_light_sigma = []
kwargs_lower_lens_light = []
kwargs_upper_lens_light = []
# first Sersic component
fixed_lens_light.append({})
kwargs_lens_light_init.append({'R_sersic': .1, 'n_sersic': 4, 'e1': 0, 'e2': 0, 'center_x': 0, 'center_y': 0})
kwargs_lens_light_sigma.append({'n_sersic': 0.5, 'R_sersic': 0.2, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1})
kwargs_lower_lens_light.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.01, 'n_sersic': 0.5, 'center_x': -10, 'center_y': -10})
kwargs_upper_lens_light.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 10, 'n_sersic': 8, 'center_x': 10, 'center_y': 10})
lens_light_params = [kwargs_lens_light_init, kwargs_lens_light_sigma, fixed_lens_light, kwargs_lower_lens_light, kwargs_upper_lens_light]
kwargs_params = {'lens_light_model': lens_light_params}
fitting_seq = FittingSequence(multi_band_list, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params)
# n_particles is the number of particles when you can test the parementers, so if you have 5 the code will be tested in 5 positions
# more particles means the fitting is tested in more places so the guess will not converge a local minimum
# but more particles is also time consuming
# n_iterations is the number of iterations in the code and there is no trivial way to know how many are necesary to make the code converge
fitting_kwargs_list = [{'fitting_routine': 'PSO', 'mpi': False, 'sigma_scale': 1., 'n_particles': 100,'n_iterations': 90}]
lens_result, source_result, lens_light_result, ps_result, cosmo_result, chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list)
lensPlot = LensModelPlot(kwargs_data_mask, kwargs_psf, kwargs_numerics, kwargs_model, lens_result, source_result,
lens_light_result, ps_result, arrow_size=0.02, cmap_string="gist_heat")
# this contain the best parameters from fitting the light of the galaxy (light, ellipticity and position)
LLR = lens_light_result[0]
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
}
kwargs_likelihood = {'check_bounds': True,
'force_no_add_image': False,
'source_marg': False,
'image_position_uncertainty': 0.004,
'check_matched_source_position': True,
'source_position_tolerance': 0.001,
'time_delay_likelihood': True,
}
kwargs_numerics = {'supersampling_factor': 1}
image_band = [kwargs_data, kwargs_psf, kwargs_numerics]
multi_band_list = [image_band]
kwargs_data_joint = {'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear',
'time_delays_measured': delta_t["delta_t"].to_numpy(),
'time_delays_uncertainties': delta_t["sigma"].to_numpy(),}
from lenstronomy.Workflow.fitting_sequence import FittingSequence
fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params)
fitting_kwargs_list = [
['PSO', {'sigma_scale': .1, 'n_particles': 200, 'n_iterations': 200}],
['MCMC', {'n_burn': 100, 'n_run': 100, 'walkerRatio': 10, 'sigma_scale': .1}]
]
start_time = time.time()
chain_list = fitting_seq.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_seq.best_fit()
end_time = time.time()
print(end_time - start_time, 'total time needed for computation')
print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ')
def 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_qso_multiband(QSO_im_list, psf_ave_list, psf_std_list=None, source_params=None,ps_param=None,
background_rms_list=[0.04]*5, pix_sz = 0.168,
exp_time = 300., fix_n=None, image_plot = True, corner_plot=True,
flux_ratio_plot=True, deep_seed = False, fixcenter = False, QSO_msk_list=None,
QSO_std_list=None, tag = None, no_MCMC= False, pltshow = 1, new_band_seq=None):
'''
A quick fit for the QSO image with (so far) single sersice + one PSF. The input psf noise is optional.
Parameter
--------
QSO_im: An array of the QSO image.
psf_ave: The psf image.
psf_std: The psf noise, optional.
source_params: The prior for the source. Default is given.
background_rms: default as 0.04
exp_time: default at 2400.
deep_seed: if Ture, more mcmc steps will be performed.
tag: The name tag for save the plot
Return
--------
Will output the fitted image (Set image_plot = True), the corner_plot and the flux_ratio_plot.
source_result, ps_result, image_ps, image_host
To do
--------
'''
# data specifics need to set up based on the data situation
background_rms_list = background_rms_list # background noise per pixel (Gaussian)
exp_time = exp_time # exposure time (arbitrary units, flux per pixel is in units #photons/exp_time unit)
numPix = len(QSO_im_list[0]) # cutout pixel size
deltaPix = pix_sz
psf_type = 'PIXEL' # 'gaussian', 'pixel', 'NONE'
kernel_list = psf_ave_list
if new_band_seq == None:
new_band_seq= range(len(QSO_im_list))
# if psf_std_list is not None:
# kwargs_numerics_list = [{'subgrid_res': 1, 'psf_subgrid': False, 'psf_error_map': True}] * len(QSO_im_list) #Turn on the PSF error map
# else:
kwargs_numerics_list = [{'supersampling_factor': 1, 'supersampling_convolution': False}] * len(QSO_im_list)
if source_params is None:
# here are the options for the host galaxy fitting
fixed_source = []
kwargs_source_init = []
kwargs_source_sigma = []
kwargs_lower_source = []
kwargs_upper_source = []
# Disk component, as modelled by an elliptical Sersic profile
if fix_n == None:
fixed_source.append({}) # we fix the Sersic index to n=1 (exponential)
kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': 2., 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.})
kwargs_source_sigma.append({'n_sersic': 0.5, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1})
kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': 0.3, 'center_x': -10, 'center_y': -10})
kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3., 'n_sersic': 7., 'center_x': 10, 'center_y': 10})
elif fix_n is not None:
fixed_source.append({'n_sersic': fix_n})
kwargs_source_init.append({'R_sersic': 0.3, 'n_sersic': fix_n, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0.})
kwargs_source_sigma.append({'n_sersic': 0.001, 'R_sersic': 0.5, 'e1': 0.1, 'e2': 0.1, 'center_x': 0.1, 'center_y': 0.1})
kwargs_lower_source.append({'e1': -0.5, 'e2': -0.5, 'R_sersic': 0.1, 'n_sersic': fix_n, 'center_x': -10, 'center_y': -10})
kwargs_upper_source.append({'e1': 0.5, 'e2': 0.5, 'R_sersic': 3, 'n_sersic': fix_n, 'center_x': 10, 'center_y': 10})
source_params = [kwargs_source_init, kwargs_source_sigma, fixed_source, kwargs_lower_source, kwargs_upper_source]
else:
source_params = source_params
if ps_param is None:
center_x = 0.0
center_y = 0.0
point_amp = QSO_im_list[0].sum()/2.
fixed_ps = [{}]
kwargs_ps = [{'ra_image': [center_x], 'dec_image': [center_y], 'point_amp': [point_amp]}]
kwargs_ps_init = kwargs_ps
kwargs_ps_sigma = [{'ra_image': [0.01], 'dec_image': [0.01]}]
kwargs_lower_ps = [{'ra_image': [-10], 'dec_image': [-10]}]
kwargs_upper_ps = [{'ra_image': [10], 'dec_image': [10]}]
ps_param = [kwargs_ps_init, kwargs_ps_sigma, fixed_ps, kwargs_lower_ps, kwargs_upper_ps]
else:
ps_param = ps_param
kwargs_params = {'source_model': source_params,
'point_source_model': ps_param}
#==============================================================================
#Doing the QSO fitting
#==============================================================================
kwargs_data_list, data_class_list = [], []
for i in range(len(QSO_im_list)):
kwargs_data_i = sim_util.data_configure_simple(numPix, deltaPix, exp_time, background_rms_list[i], inverse=True)
kwargs_data_list.append(kwargs_data_i)
data_class_list.append(ImageData(**kwargs_data_i))
kwargs_psf_list = []
psf_class_list = []
for i in range(len(QSO_im_list)):
kwargs_psf_i = {'psf_type': psf_type, 'kernel_point_source': kernel_list[i]}
kwargs_psf_list.append(kwargs_psf_i)
psf_class_list.append(PSF(**kwargs_psf_i))
data_class_list[i].update_data(QSO_im_list[i])
light_model_list = ['SERSIC_ELLIPSE'] * len(source_params[0])
lightModel = LightModel(light_model_list=light_model_list)
point_source_list = ['UNLENSED']
pointSource = PointSource(point_source_type_list=point_source_list)
imageModel_list = []
for i in range(len(QSO_im_list)):
kwargs_data_list[i]['image_data'] = QSO_im_list[i]
# if QSO_msk_list is not None:
# kwargs_numerics_list[i]['mask'] = QSO_msk_list[i]
if QSO_std_list is not None:
kwargs_data_list[i]['noise_map'] = QSO_std_list[i]
# if psf_std_list is not None:
# kwargs_psf_list[i]['psf_error_map'] = psf_std_list[i]
image_band_list = []
for i in range(len(QSO_im_list)):
imageModel_list.append(ImageModel(data_class_list[i], psf_class_list[i], source_model_class=lightModel,
point_source_class=pointSource, kwargs_numerics=kwargs_numerics_list[i]))
image_band_list.append([kwargs_data_list[i], kwargs_psf_list[i], kwargs_numerics_list[i]])
multi_band_list = [image_band_list[i] for i in range(len(QSO_im_list))]
# numerical options and fitting sequences
kwargs_model = { 'source_light_model_list': light_model_list,
'point_source_model_list': point_source_list
}
if fixcenter == False:
kwargs_constraints = {'num_point_source_list': [1]
}
elif fixcenter == True:
kwargs_constraints = {'joint_source_with_point_source': [[0, 0]],
'num_point_source_list': [1]
}
kwargs_likelihood = {'check_bounds': True, #Set the bonds, if exceed, reutrn "penalty"
'source_marg': False, #In likelihood_module.LikelihoodModule -- whether to fully invert the covariance matrix for marginalization
'check_positive_flux': True,
'image_likelihood_mask_list': [QSO_msk_list]
}
# mpi = False # MPI possible, but not supported through that notebook.
# The Params for the fitting. kwargs_init: initial input. kwargs_sigma: The parameter uncertainty. kwargs_fixed: fixed parameters;
#kwargs_lower,kwargs_upper: Lower and upper limits.
kwargs_data_joint = {'multi_band_list': multi_band_list, 'multi_band_type': 'multi-linear'} # 'single-band', 'multi-linear', 'joint-linear'
fitting_seq = FittingSequence(kwargs_data_joint, kwargs_model, kwargs_constraints, kwargs_likelihood, kwargs_params)
if deep_seed == False:
fitting_kwargs_list = [
['PSO', {'sigma_scale': 0.8, 'n_particles': 80, 'n_iterations': 60, 'compute_bands': [True]+[False]*(len(QSO_im_list)-1)}],
['align_images', {'n_particles': 10, 'n_iterations': 10, 'compute_bands': [False]+[True]*(len(QSO_im_list)-1)}],
['PSO', {'sigma_scale': 0.8, 'n_particles': 100, 'n_iterations': 200, 'compute_bands': [True]*len(QSO_im_list)}],
['MCMC', {'n_burn': 10, 'n_run': 20, 'walkerRatio': 50, 'sigma_scale': .1}]
]
elif deep_seed == True:
fitting_kwargs_list = [
['PSO', {'sigma_scale': 0.8, 'n_particles': 150, 'n_iterations': 60, 'compute_bands': [True]+[False]*(len(QSO_im_list)-1)}],
['align_images', {'n_particles': 20, 'n_iterations': 20, 'compute_bands': [False]+[True]*(len(QSO_im_list)-1)}],
['PSO', {'sigma_scale': 0.8, 'n_particles': 150, 'n_iterations': 200, 'compute_bands': [True]*len(QSO_im_list)}],
['MCMC', {'n_burn': 20, 'n_run': 40, 'walkerRatio': 50, 'sigma_scale': .1}]
]
if no_MCMC == True:
del fitting_kwargs_list[-1]
start_time = time.time()
# lens_result, source_result, lens_light_result, ps_result, cosmo_temp, chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list)
chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list)
lens_result, source_result, lens_light_result, ps_result, cosmo_temp = fitting_seq.best_fit()
end_time = time.time()
print(end_time - start_time, 'total time needed for computation')
print('============ CONGRATULATION, YOUR JOB WAS SUCCESSFUL ================ ')
source_result_list, ps_result_list = [], []
image_reconstructed_list, error_map_list, image_ps_list, image_host_list, shift_RADEC_list=[], [], [], [],[]
imageLinearFit_list = []
for k in range(len(QSO_im_list)):
# this is the linear inversion. The kwargs will be updated afterwards
imageLinearFit_k = ImageLinearFit(data_class_list[k], psf_class_list[k], source_model_class=lightModel,
point_source_class=pointSource, kwargs_numerics=kwargs_numerics_list[k])
image_reconstructed_k, error_map_k, _, _ = imageLinearFit_k.image_linear_solve(kwargs_source=source_result, kwargs_ps=ps_result)
imageLinearFit_list.append(imageLinearFit_k)
[kwargs_data_k, kwargs_psf_k, kwargs_numerics_k] = fitting_seq.multi_band_list[k]
# data_class_k = data_class_list[k] #ImageData(**kwargs_data_k)
# psf_class_k = psf_class_list[k] #PSF(**kwargs_psf_k)
# imageModel_k = ImageModel(data_class_k, psf_class_k, source_model_class=lightModel,
# point_source_class=pointSource, kwargs_numerics=kwargs_numerics_list[k])
imageModel_k = imageModel_list[k]
modelPlot = ModelPlot(multi_band_list[k], kwargs_model, lens_result, source_result,
lens_light_result, ps_result, arrow_size=0.02, cmap_string="gist_heat", likelihood_mask=QSO_im_list[k])
print("source_result", 'for', "k", source_result)
image_host_k = []
for i in range(len(source_result)):
image_host_k.append(imageModel_list[k].source_surface_brightness(source_result,de_lensed=True,unconvolved=False, k=i))
image_ps_k = imageModel_k.point_source(ps_result)
# let's plot the output of the PSO minimizer
image_reconstructed_list.append(image_reconstructed_k)
source_result_list.append(source_result)
ps_result_list.append(ps_result)
error_map_list.append(error_map_k)
image_ps_list.append(image_ps_k)
image_host_list.append(image_host_k)
if 'ra_shift' in fitting_seq.multi_band_list[k][0].keys():
shift_RADEC_list.append([fitting_seq.multi_band_list[k][0]['ra_shift'], fitting_seq.multi_band_list[k][0]['dec_shift']])
else:
shift_RADEC_list.append([0,0])
if image_plot:
f, axes = plt.subplots(3, 3, figsize=(16, 16), sharex=False, sharey=False)
modelPlot.data_plot(ax=axes[0,0], text="Data")
modelPlot.model_plot(ax=axes[0,1])
modelPlot.normalized_residual_plot(ax=axes[0,2], v_min=-6, v_max=6)
modelPlot.decomposition_plot(ax=axes[1,0], text='Host galaxy', source_add=True, unconvolved=True)
modelPlot.decomposition_plot(ax=axes[1,1], text='Host galaxy convolved', source_add=True)
modelPlot.decomposition_plot(ax=axes[1,2], text='All components convolved', source_add=True, lens_light_add=True, point_source_add=True)
modelPlot.subtract_from_data_plot(ax=axes[2,0], text='Data - Point Source', point_source_add=True)
modelPlot.subtract_from_data_plot(ax=axes[2,1], text='Data - host galaxy', source_add=True)
modelPlot.subtract_from_data_plot(ax=axes[2,2], text='Data - host galaxy - Point Source', source_add=True, point_source_add=True)
f.tight_layout()
if tag is not None:
f.savefig('{0}_fitted_image_band{1}.pdf'.format(tag,new_band_seq[k]))
if pltshow == 0:
plt.close()
else:
plt.show()
if corner_plot==True and no_MCMC==False and k ==0:
# here the (non-converged) MCMC chain of the non-linear parameters
if not samples_mcmc == []:
n, num_param = np.shape(samples_mcmc)
plot = corner.corner(samples_mcmc, labels=param_mcmc, show_titles=True)
if tag is not None:
plot.savefig('{0}_para_corner.pdf'.format(tag))
if pltshow == 0:
plt.close()
else:
plt.show()
if flux_ratio_plot==True and no_MCMC==False:
param = Param(kwargs_model, kwargs_fixed_source=source_params[2], kwargs_fixed_ps=fixed_ps, **kwargs_constraints)
mcmc_new_list = []
labels_new = [r"Quasar flux", r"host_flux", r"source_x", r"source_y"]
# transform the parameter position of the MCMC chain in a lenstronomy convention with keyword arguments #
for i in range(len(samples_mcmc)/10):
kwargs_lens_out, kwargs_light_source_out, kwargs_light_lens_out, kwargs_ps_out, kwargs_cosmo = param.getParams(samples_mcmc[i+ len(samples_mcmc)/10*9])
image_reconstructed, _, _, _ = imageLinearFit_list[k].image_linear_solve(kwargs_source=kwargs_light_source_out, kwargs_ps=kwargs_ps_out)
image_ps = imageModel_list[k].point_source(kwargs_ps_out)
flux_quasar = np.sum(image_ps)
image_disk = imageModel_list[k].source_surface_brightness(kwargs_light_source_out,de_lensed=True,unconvolved=False, k=0)
flux_disk = np.sum(image_disk)
source_x = kwargs_ps_out[0]['ra_image']
source_y = kwargs_ps_out[0]['dec_image']
if flux_disk>0:
mcmc_new_list.append([flux_quasar, flux_disk, source_x, source_y])
plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True)
if tag is not None:
plot.savefig('{0}_HOSTvsQSO_corner_band{1}.pdf'.format(tag,new_band_seq[k]))
if pltshow == 0:
plt.close()
else:
plt.show()
errp_list = []
for k in range(len(QSO_im_list)):
if QSO_std_list is None:
errp_list.append(np.sqrt(data_class_list[k].C_D+np.abs(error_map_list[k])))
else:
errp_list.append(np.sqrt(QSO_std_list[k]**2+np.abs(error_map_list[k])))
return source_result_list, ps_result_list, image_ps_list, image_host_list, errp_list, shift_RADEC_list, fitting_seq #fitting_seq.multi_band_list
def sl_sys_analysis():
# Get command line arguments
args = {}
if comm_rank == 0:
print(":Registered %d processes" % comm_size)
args["infile"] = sys.argv[1]
args["nimgs"] = sys.argv[2]
args["los"] = sys.argv[3]
args["version"] = sys.argv[4]
args["dt_sigma"] = float(sys.argv[5])
args["image_amps_sigma"] = float(sys.argv[6])
args["flux_ratio_errors"] = float(sys.argv[7])
args["astrometry_sigma"] = float(sys.argv[8])
args = comm.bcast(args)
# Organize devision of strong lensing systems
with open(args["infile"], "r") as myfile:
limg_data = myfile.read()
systems = json.loads(limg_data)
sys_nr_per_proc = int(len(systems) / comm_size)
print("comm_rank", comm_rank)
start_sys = sys_nr_per_proc * comm_rank
end_sys = sys_nr_per_proc * (comm_rank + 1)
print(start_sys, end_sys)
with open("../lens_catalogs_sie_only.json", "r") as myfile:
limg_data = myfile.read()
systems_prior = json.loads(limg_data)
if comm_rank == 0:
print("Each process will have %d systems" % sys_nr_per_proc)
print("That should take app. %f min." % (sys_nr_per_proc * 20))
source_size_pc = 10.0
window_size = 0.1 # units of arcseconds
grid_number = 100 # supersampled window (per axis)
z_source = 2.0
cosmo = FlatLambdaCDM(H0=71, Om0=0.3089, Ob0=0.0)
results = {"gamma": [], "phi_ext": [], "gamma_ext": [], "theta_E": [], "D_dt": []}
for ii in range(len(systems))[(start_sys + 2) : end_sys]:
system = systems[ii]
system_prior = systems_prior[ii]
print("Analysing system ID: %d" % ii)
# the data set is
z_lens = system_prior["zl"]
lensCosmo = LensCosmo(cosmo=cosmo, z_lens=z_lens, z_source=z_source)
# convert units of pc into arcseconds
D_s = lensCosmo.D_s
source_size_arcsec = source_size_pc / 10 ** 6 / D_s / constants.arcsec
print("The source size in arcsec init = %.4f" % source_size_arcsec) #0.0012
# multiple images properties
ximg = np.zeros(system["nimgs"])
yimg = np.zeros(system["nimgs"])
t_days = np.zeros(system["nimgs"])
image_amps = np.zeros(system["nimgs"])
for jj in range(system["nimgs"]):
ximg[jj] = system["ximg"][jj] # [arcsec]
yimg[jj] = system["yimg"][jj] # [arcsec]
t_days[jj] = system["delay"][jj] # [days]
image_amps[jj] = system["mags"][jj] # [linear units or magnitudes]
# sort by arrival time
index_sort = np.argsort(t_days)
ximg = ximg[index_sort] # relative RA (arc seconds)
yimg = yimg[index_sort] # relative DEC (arc seconds)
image_amps = np.abs(image_amps[index_sort])
t_days = t_days[index_sort]
d_dt = t_days[1:] - t_days[0]
# measurement uncertainties
astrometry_sigma = args["astrometry_sigma"]
ximg_measured = ximg + np.random.normal(0, astrometry_sigma, system["nimgs"])
yimg_measured = yimg + np.random.normal(0, astrometry_sigma, system["nimgs"])
image_amps_sigma = np.ones(system["nimgs"]) * args["image_amps_sigma"]
flux_ratios = image_amps[1:] - image_amps[0]
flux_ratio_errors = np.ones(system["nimgs"] - 1) * args["flux_ratio_errors"]
flux_ratios_measured = flux_ratios + np.random.normal(0, flux_ratio_errors)
d_dt_sigma = np.ones(system["nimgs"] - 1) * args["dt_sigma"]
d_dt_measured = d_dt + np.random.normal(0, d_dt_sigma)
kwargs_data_joint = {
"time_delays_measured": d_dt_measured,
"time_delays_uncertainties": d_dt_sigma,
"flux_ratios": flux_ratios_measured,
"flux_ratio_errors": flux_ratio_errors,
"ra_image_list": [ximg_measured],
"dec_image_list": [yimg_measured],
}
# lens model choices
lens_model_list = ["SPEMD", "SHEAR_GAMMA_PSI"]
# 1. layer: primary SPEP
fixed_lens = []
kwargs_lens_init = []
kwargs_lens_sigma = []
kwargs_lower_lens = []
kwargs_upper_lens = []
fixed_lens.append({})
kwargs_lens_init.append(
{
"theta_E": 1.0,
"gamma": 2,
"center_x": 0,
"center_y": 0,
"e1": 0,
"e2": 0.0,
}
)
# error
kwargs_lens_sigma.append(
{
"theta_E": 0.2,
"e1": 0.1,
"e2": 0.1,
"gamma": 0.1,
"center_x": 0.1,
"center_y": 0.1,
}
)
# lower limit
kwargs_lower_lens.append(
{
"theta_E": 0.01,
"e1": -0.5,
"e2": -0.5,
"gamma": 1.5,
"center_x": -10,
"center_y": -10,
}
)
# upper limit
kwargs_upper_lens.append(
{
"theta_E": 10,
"e1": 0.5,
"e2": 0.5,
"gamma": 2.5,
"center_x": 10,
"center_y": 10,
}
)
# 2nd layer: external SHEAR
fixed_lens.append({"ra_0": 0, "dec_0": 0})
kwargs_lens_init.append({"gamma_ext": 0.05, "psi_ext": 0.0})
kwargs_lens_sigma.append({"gamma_ext": 0.05, "psi_ext": np.pi})
kwargs_lower_lens.append({"gamma_ext": 0, "psi_ext": -np.pi})
kwargs_upper_lens.append({"gamma_ext": 0.3, "psi_ext": np.pi})
# 3rd layer: external CONVERGENCE
kwargs_lens_init.append({'kappa_ext': 0.12})
kwargs_lens_sigma.append({'kappa_ext': 0.06})
kwargs_lower_lens.append({'kappa_ext': 0.0})
kwargs_upper_lens.append({'kappa_ext': 0.3})
# combined lens model
lens_params = [
kwargs_lens_init,
kwargs_lens_sigma,
fixed_lens,
kwargs_lower_lens,
kwargs_upper_lens,
]
# image position parameters
point_source_list = ["LENSED_POSITION"]
# we fix the image position coordinates
fixed_ps = [{}]
# the initial guess for the appearing image positions is:
# at the image position.
kwargs_ps_init = [{"ra_image": ximg, "dec_image": yimg}]
# let some freedome in how well the actual image positions are
# matching those given by the data (indicated as 'ra_image', 'dec_image'
# and held fixed while fitting)
kwargs_ps_sigma = [
{
"ra_image": 0.01 * np.ones(len(ximg)),
"dec_image": 0.01 * np.ones(len(ximg)),
}
]
kwargs_lower_ps = [
{
"ra_image": -10 * np.ones(len(ximg)),
"dec_image": -10 * np.ones(len(ximg)),
}
]
kwargs_upper_ps = [
{"ra_image": 10 * np.ones(len(ximg)), "dec_image": 10 * np.ones(len(ximg))}
]
ps_params = [
kwargs_ps_init,
kwargs_ps_sigma,
fixed_ps,
kwargs_lower_ps,
kwargs_upper_ps,
]
# quasar source size
fixed_special = {}
kwargs_special_init = {}
kwargs_special_sigma = {}
kwargs_lower_special = {}
kwargs_upper_special = {}
fixed_special["source_size"] = source_size_arcsec
kwargs_special_init["source_size"] = source_size_arcsec
kwargs_special_sigma["source_size"] = source_size_arcsec
kwargs_lower_special["source_size"] = 0.0001
kwargs_upper_special["source_size"] = 1
# Time-delay distance
kwargs_special_init["D_dt"] = 4300 # corresponds to H0 ~ 70
kwargs_special_sigma["D_dt"] = 3000
kwargs_lower_special["D_dt"] = 2500 # corresponds to H0 ~ 120
kwargs_upper_special["D_dt"] = 14000 # corresponds to H0 ~ 20
special_params = [
kwargs_special_init,
kwargs_special_sigma,
fixed_special,
kwargs_lower_special,
kwargs_upper_special,
]
# combined parameter settings
kwargs_params = {
"lens_model": lens_params,
"point_source_model": ps_params,
"special": special_params,
}
# our model choices
kwargs_model = {
"lens_model_list": lens_model_list,
"point_source_model_list": point_source_list,
}
lensModel = LensModel(kwargs_model["lens_model_list"])
lensModelExtensions = LensModelExtensions(lensModel=lensModel)
lensEquationSolver = LensEquationSolver(lensModel=lensModel)
# setup options for likelihood and parameter sampling
time_delay_likelihood = True
flux_ratio_likelihood = True
image_position_likelihood = True
kwargs_flux_compute = {
"source_type": "INF",
"window_size": window_size,
"grid_number": grid_number,
}
kwargs_constraints = {
"num_point_source_list": [int(args["nimgs"])],
# any proposed lens model must satisfy the image positions
# appearing at the position of the point sources being sampeld
# "solver_type": "PROFILE_SHEAR",
"Ddt_sampling": time_delay_likelihood,
# sampling of the time-delay distance
# explicit modelling of the astrometric imperfection of
# the point source positions
"point_source_offset": True,
}
# explicit sampling of finite source size parameter
# (only use when source_type='GAUSSIAN' or 'TORUS')
if (
kwargs_flux_compute["source_type"] in ["GAUSSIAN", "TORUS"]
and flux_ratio_likelihood is True
):
kwargs_constraints["source_size"] = True
# e.g. power-law mass slope of the main deflector
# [[index_model, 'param_name', mean, 1-sigma error], [...], ...]
prior_lens = [[0, "gamma", 2, 0.1]]
prior_special = []
kwargs_likelihood = {
"position_uncertainty": args["astrometry_sigma"],
"source_position_likelihood": True,
"image_position_likelihood": True,
"time_delay_likelihood": True,
"flux_ratio_likelihood": True,
"kwargs_flux_compute": kwargs_flux_compute,
"prior_lens": prior_lens,
"prior_special": prior_special,
"check_solver": True,
"solver_tolerance": 0.001,
"check_bounds": True,
}
fitting_seq = FittingSequence(
kwargs_data_joint,
kwargs_model,
kwargs_constraints,
kwargs_likelihood,
kwargs_params,
)
fitting_kwargs_list = [
["PSO", {"sigma_scale": 1.0, "n_particles": 200, "n_iterations": 500}]
]
chain_list_pso = fitting_seq.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_seq.best_fit()
kwargs_result = fitting_seq.best_fit(bijective=True)
args_result = fitting_seq.param_class.kwargs2args(**kwargs_result)
logL, _ = fitting_seq.likelihoodModule.logL(args_result, verbose=True)
# and now we run the MCMC
fitting_kwargs_list = [
[
"MCMC",
{"n_burn": 400, "n_run": 600, "walkerRatio": 10, "sigma_scale": 0.1},
]
]
chain_list_mcmc = fitting_seq.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_seq.best_fit()
# print("number of non-linear parameters in the MCMC process: ", len(param_mcmc))
# print("parameters in order: ", param_mcmc)
print("number of evaluations in the MCMC process: ", np.shape(samples_mcmc)[0])
param = Param(
kwargs_model,
fixed_lens,
kwargs_fixed_ps=fixed_ps,
kwargs_fixed_special=fixed_special,
kwargs_lens_init=kwargs_result["kwargs_lens"],
**kwargs_constraints,
)
# the number of non-linear parameters and their names #
num_param, param_list = param.num_param()
for i in range(len(samples_mcmc)):
kwargs_out = param.args2kwargs(samples_mcmc[i])
kwargs_lens_out, kwargs_special_out, kwargs_ps_out = (
kwargs_out["kwargs_lens"],
kwargs_out["kwargs_special"],
kwargs_out["kwargs_ps"],
)
# compute 'real' image position adding potential astrometric shifts
x_pos = kwargs_ps_out[0]["ra_image"]
y_pos = kwargs_ps_out[0]["dec_image"]
# extract quantities of the main deflector
theta_E = kwargs_lens_out[0]["theta_E"]
gamma = kwargs_lens_out[0]["gamma"]
e1, e2 = kwargs_lens_out[0]["e1"], kwargs_lens_out[0]["e2"]
phi, q = param_util.ellipticity2phi_q(e1, e2)
phi_ext, gamma_ext = (
kwargs_lens_out[1]["psi_ext"] % np.pi,
kwargs_lens_out[1]["gamma_ext"],
)
if flux_ratio_likelihood is True:
mag = lensModel.magnification(x_pos, y_pos, kwargs_lens_out)
flux_ratio_fit = mag[1:] / mag[0]
if (
kwargs_constraints.get("source_size", False) is True
and "source_size" not in fixed_special
):
source_size = kwargs_special_out["source_size"]
if time_delay_likelihood is True:
D_dt = kwargs_special_out["D_dt"]
# and here the predicted angular diameter distance from a
# default cosmology (attention for experimenter bias!)
gamma = np.median(gamma)
phi_ext = np.median(phi_ext)
gamma_ext = np.median(gamma_ext)
theta_E = np.median(theta_E)
D_dt = np.median(D_dt)
results["gamma"].append(gamma)
results["phi_ext"].append(phi_ext)
results["gamma_ext"].append(gamma_ext)
results["theta_E"].append(theta_E)
results["H0"].append(c_light / D_dt)
def sl_sys_analysis():
# Get command line arguments
args = {}
if comm_rank == 0:
print(":Registered %d processes" % comm_size)
args["infile"] = sys.argv[1]
args["nimgs"] = sys.argv[2]
args["los"] = sys.argv[3]
args["version"] = sys.argv[4]
args["dt_sigma"] = float(sys.argv[5])
args["image_amps_sigma"] = float(sys.argv[6])
args["flux_ratio_errors"] = float(sys.argv[7])
args["astrometry_sigma"] = float(sys.argv[8])
args = comm.bcast(args)
# Organize devision of strong lensing systems
with open(args["infile"], "r") as myfile:
limg_data = myfile.read()
systems = json.loads(limg_data)
sys_nr_per_proc = int(len(systems) / comm_size)
print('comm_rank', comm_rank)
start_sys = sys_nr_per_proc * comm_rank
end_sys = sys_nr_per_proc * (comm_rank + 1)
print(start_sys, end_sys)
with open("../lens_catalogs_sie_only.json", "r") as myfile:
limg_data = myfile.read()
systems_prior = json.loads(limg_data)
if comm_rank == 0:
print("Each process will have %d systems" % sys_nr_per_proc)
print("That should take app. %f min." % (sys_nr_per_proc * 20))
results = {
"gamma": [],
"phi_ext": [],
"gamma_ext": [],
"theta_E": [],
"D_dt": [],
}
for ii in range(len(systems))[(start_sys + 2):end_sys]:
system = systems[ii]
system_prior = systems_prior[ii]
print("Analysing system ID: %d" % ii)
print(system)
# the data set is
z_lens = system_prior["zl"]
z_source = 2.0
# multiple images properties
ximg = np.zeros(system["nimgs"])
yimg = np.zeros(system["nimgs"])
delay = np.zeros(system["nimgs"])
image_amps = np.zeros(system["nimgs"])
for jj in range(system["nimgs"]):
ximg[jj] = system["ximg"][jj] #[arcsec]
yimg[jj] = system["yimg"][jj] #[arcsec]
delay[jj] = system["delay"][jj] #[days]
image_amps[jj] = system["mags"][jj] #[linear units or magnitudes]
# sort by arrival time
index_sort = np.argsort(delay)
ximg = ximg[index_sort]
yimg = yimg[index_sort]
delay = delay[index_sort]
image_amps = image_amps[index_sort]
d_dt = delay[1:] - delay[0]
# measurement uncertainties
d_dt_sigma = np.ones(system["nimgs"] - 1) * args["dt_sigma"]
image_amps_sigma = np.ones(system["nimgs"]) * args["image_amps_sigma"]
flux_ratios = image_amps[1:] - image_amps[0]
flux_ratio_errors = np.ones(system["nimgs"] -
1) * args["flux_ratio_errors"]
# lens model choices
lens_model_list = ["SPEP", "SHEAR"]
# first choice: SPEP
fixed_lens = []
kwargs_lens_init = []
kwargs_lens_sigma = []
kwargs_lower_lens = []
kwargs_upper_lens = []
fixed_lens.append({})
kwargs_lens_init.append({
"theta_E": 1.0,
"gamma": 2,
"center_x": 0,
"center_y": 0,
"e1": 0,
"e2": 0.0,
})
# error
kwargs_lens_sigma.append({
"theta_E": 0.2,
"e1": 0.1,
"e2": 0.1,
"gamma": 0.1,
"center_x": 0.1,
"center_y": 0.1,
})
# lower limit
kwargs_lower_lens.append({
"theta_E": 0.01,
"e1": -0.5,
"e2": -0.5,
"gamma": 1.5,
"center_x": -10,
"center_y": -10,
})
# upper limit
kwargs_upper_lens.append({
"theta_E": 10,
"e1": 0.5,
"e2": 0.5,
"gamma": 2.5,
"center_x": 10,
"center_y": 10,
})
# second choice: SHEAR
fixed_lens.append({"ra_0": 0, "dec_0": 0})
kwargs_lens_init.append({"e1": 0.0, "e2": 0.0})
kwargs_lens_sigma.append({"e1": 0.1, "e2": 0.1})
kwargs_lower_lens.append({"e1": -0.2, "e2": -0.2})
kwargs_upper_lens.append({"e1": 0.2, "e2": 0.2})
lens_params = [
kwargs_lens_init,
kwargs_lens_sigma,
fixed_lens,
kwargs_lower_lens,
kwargs_upper_lens,
]
point_source_list = ["LENSED_POSITION"]
fixed_ps = [{"ra_image": ximg, "dec_image": yimg}]
kwargs_ps_init = fixed_ps
# let some freedome in how well the actual image positions are
# matching those given by the data (indicated as 'ra_image', 'dec_image'
# and held fixed while fitting)
kwargs_ps_sigma = [{
"ra_image": 0.01 * np.ones(len(ximg)),
"dec_image": 0.01 * np.ones(len(ximg)),
}]
kwargs_lower_ps = [{
"ra_image": -10 * np.ones(len(ximg)),
"dec_image": -10 * np.ones(len(ximg)),
}]
kwargs_upper_ps = [{
"ra_image": 10 * np.ones(len(ximg)),
"dec_image": 10 * np.ones(len(ximg))
}]
ps_params = [
kwargs_ps_init,
kwargs_ps_sigma,
fixed_ps,
kwargs_lower_ps,
kwargs_upper_ps,
]
fixed_cosmo = {}
kwargs_cosmo_init = {
"D_dt": 5000,
"delta_x_image": np.zeros_like(ximg),
"delta_y_image": np.zeros_like(ximg),
}
kwargs_cosmo_sigma = {
"D_dt": 10000,
"delta_x_image": np.ones_like(ximg) * args["astrometry_sigma"],
"delta_y_image": np.ones_like(ximg) * args["astrometry_sigma"],
}
kwargs_lower_cosmo = {
"D_dt": 0,
"delta_x_image": np.ones_like(ximg) * (-1),
"delta_y_image": np.ones_like(ximg) * (-1),
}
kwargs_upper_cosmo = {
"D_dt": 10000,
"delta_x_image": np.ones_like(ximg) * (1),
"delta_y_image": np.ones_like(ximg) * (1),
}
cosmo_params = [
kwargs_cosmo_init,
kwargs_cosmo_sigma,
fixed_cosmo,
kwargs_lower_cosmo,
kwargs_upper_cosmo,
]
kwargs_params = {
"lens_model": lens_params,
"point_source_model": ps_params,
"cosmography": cosmo_params,
}
# setup options for likelihood and parameter sampling
kwargs_constraints = {
"num_point_source_list": [int(args["nimgs"])],
# any proposed lens model must satisfy the image positions
# appearing at the position of the point sources being sampeld
"solver_type": "PROFILE_SHEAR",
"cosmo_type": "D_dt",
# sampling of the time-delay distance
# explicit modelling of the astrometric imperfection of
# the point source positions
"point_source_offset": True,
}
kwargs_likelihood = {
"check_bounds": True,
"point_source_likelihood": True,
"position_uncertainty": args["astrometry_sigma"],
"check_solver": True,
"solver_tolerance": 0.001,
"time_delay_likelihood": True,
"image_likelihood": False,
# this needs to be explicitly given when not having imaging data
"flux_ratio_likelihood": True, # enables the flux ratio likelihood
}
kwargs_data_joint = {
"time_delays_measured": d_dt,
"time_delays_uncertainties": d_dt_sigma,
"flux_ratios": flux_ratios,
"flux_ratio_errors": flux_ratio_errors,
}
kwargs_model = {
"lens_model_list": lens_model_list,
"point_source_model_list": point_source_list,
}
mpi = False # MPI possible, but not supported through that notebook.
fitting_seq = FittingSequence(
kwargs_data_joint,
kwargs_model,
kwargs_constraints,
kwargs_likelihood,
kwargs_params,
)
fitting_kwargs_list = [
# ['update_settings', {'lens_add_fixed': [[0, ['gamma']]]}],
[
"PSO", {
"sigma_scale": 1.0,
"n_particles": 100,
"n_iterations": 100
}
]
]
chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(
fitting_kwargs_list)
lens_result, source_result, lens_light_result, ps_result, cosmo_result = (
fitting_seq.best_fit())
# and now we run the MCMC
fitting_kwargs_list = [
[
"PSO", {
"sigma_scale": 0.1,
"n_particles": 100,
"n_iterations": 100
}
],
[
"MCMC",
{
"n_burn": 200,
"n_run": 200,
"walkerRatio": 10,
"sigma_scale": 0.1
},
],
]
chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(
fitting_kwargs_list)
lens_result, source_result, lens_light_result, ps_result, cosmo_result = (
fitting_seq.best_fit())
# print("number of non-linear parameters in the MCMC process: ", len(param_mcmc))
# print("parameters in order: ", param_mcmc)
print("number of evaluations in the MCMC process: ",
np.shape(samples_mcmc)[0])
param = Param(
kwargs_model,
fixed_lens,
kwargs_fixed_ps=fixed_ps,
kwargs_fixed_cosmo=fixed_cosmo,
kwargs_lens_init=lens_result,
**kwargs_constraints,
)
# the number of non-linear parameters and their names
num_param, param_list = param.num_param()
lensAnalysis = LensAnalysis(kwargs_model)
mcmc_new_list = []
labels_new = [
r"$\gamma$",
r"$\phi_{ext}$",
r"$\gamma_{ext}$",
r"$D_{\Delta t}$",
]
D_dt = np.zeros(len(samples_mcmc))
theta_E = np.zeros(len(samples_mcmc))
for i in range(len(samples_mcmc)):
# transform the parameter position of the MCMC chain in a
# lenstronomy convention with keyword arguments
kwargs_lens_out, kwargs_light_source_out, kwargs_light_lens_out, kwargs_ps_out, kwargs_cosmo = param.args2kwargs(
samples_mcmc[i])
D_dt[i] = kwargs_cosmo["D_dt"]
gamma[i] = kwargs_lens_out[0]["gamma"]
theta_E[i] = kwargs_lens_out[0]['theta_E']
e1[i] = kwargs_lens_out[0]['e1']
e2[i] = kwargs_lens_out[0]['e2']
phi_ext, gamma_ext = lensAnalysis._lensModelExtensions.external_shear(
kwargs_lens_out)
# plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True)
# and here the predicted angular diameter distance from a
# default cosmology (attention for experimenter bias!)
cosmo = FlatLambdaCDM(
H0=71,
Om0=0.3089,
Ob0=0.,
)
lensCosmo = LensCosmo(
cosmo=cosmo,
z_lens=z_lens,
z_source=z_source,
)
gamma = np.mean(gamma)
phi_ext = np.mean(phi_ext)
gamma_ext = np.mean(gamma_ext)
theta_E = np.mean(theta_E)
D_dt = np.mean(D_dt)
results["gamma"].append(gamma)
results["phi_ext"].append(phi_ext)
results["gamma_ext"].append(gamma_ext)
results["theta_E"].append(theta_E)
results["D_dt"].append(lensCosmo.D_dt)
with open(
"./quasars_%s_nimgs_%s_%s.json" %
(args["los"], args["nimgs"], args["version"]), 'w') as fout:
json.dump(results, fout)
def test_fitting_sequence(self):
# kwargs_init = [self.kwargs_lens, self.kwargs_source, self.kwargs_lens_light, self.kwargs_ps]
lens_sigma = [{
'theta_E': 0.1,
'gamma': 0.1,
'e1': 0.1,
'e2': 0.1,
'center_x': 0.1,
'center_y': 0.1
}, {
'e1': 0.1,
'e2': 0.1
}]
lens_lower = [{
'theta_E': 0.,
'gamma': 1.5,
'center_x': -2,
'center_y': -2,
'e1': -0.4,
'e2': -0.4
}, {
'e1': -0.3,
'e2': -0.3
}]
lens_upper = [{
'theta_E': 10.,
'gamma': 2.5,
'center_x': 2,
'center_y': 2,
'e1': 0.4,
'e2': 0.4
}, {
'e1': 0.3,
'e2': 0.3
}]
source_sigma = [{
'R_sersic': 0.05,
'n_sersic': 0.5,
'center_x': 0.1,
'center_y': 0.1,
'e1': 0.1,
'e2': 0.1
}]
source_lower = [{
'R_sersic': 0.01,
'n_sersic': 0.5,
'center_x': -2,
'center_y': -2,
'e1': -0.4,
'e2': -0.4
}]
source_upper = [{
'R_sersic': 10,
'n_sersic': 5.5,
'center_x': 2,
'center_y': 2,
'e1': 0.4,
'e2': 0.4
}]
lens_light_sigma = [{
'R_sersic': 0.05,
'n_sersic': 0.5,
'center_x': 0.1,
'center_y': 0.1
}]
lens_light_lower = [{
'R_sersic': 0.01,
'n_sersic': 0.5,
'center_x': -2,
'center_y': -2
}]
lens_light_upper = [{
'R_sersic': 10,
'n_sersic': 5.5,
'center_x': 2,
'center_y': 2
}]
ps_sigma = [{'ra_source': 1, 'dec_source': 1, 'point_amp': 1}]
lens_param = self.kwargs_lens, lens_sigma, [{}, {
'ra_0': 0,
'dec_0': 0
}], lens_lower, lens_upper
source_param = self.kwargs_source, source_sigma, [
{}
], source_lower, source_upper
lens_light_param = self.kwargs_lens_light, lens_light_sigma, [{
'center_x':
0
}], lens_light_lower, lens_light_upper
ps_param = self.kwargs_ps, ps_sigma, [{}
], self.kwargs_ps, self.kwargs_ps
kwargs_params = {
'lens_model': lens_param,
'source_model': source_param,
'lens_light_model': lens_light_param,
'point_source_model': ps_param,
# 'cosmography': cosmo_param
}
# kwargs_params = [kwargs_init, kwargs_sigma, kwargs_fixed, kwargs_init, kwargs_init]
image_band = [self.kwargs_data, self.kwargs_psf, self.kwargs_numerics]
multi_band_list = [image_band]
kwargs_data_joint = {
'multi_band_list': multi_band_list,
'multi_band_type': 'multi-linear'
}
fittingSequence = FittingSequence(kwargs_data_joint, self.kwargs_model,
self.kwargs_constraints,
self.kwargs_likelihood,
kwargs_params)
kwargs_result = fittingSequence.best_fit(bijective=False)
lens_temp = kwargs_result['kwargs_lens']
npt.assert_almost_equal(lens_temp[0]['theta_E'],
self.kwargs_lens[0]['theta_E'],
decimal=2)
logL = fittingSequence.best_fit_likelihood
print(logL, 'test')
#print(lens_temp, source_temp, lens_light_temp, ps_temp, cosmo_temp)
assert logL < 0
bic = fittingSequence.bic
assert bic > 0
#npt.assert_almost_equal(bic, 20000000220.29376, decimal=-4)
#npt.assert_almost_equal(logL, -10000000061.792593, decimal=-4)
n_p = 2
n_i = 2
fitting_list = []
kwargs_pso = {
'sigma_scale': 1,
'n_particles': n_p,
'n_iterations': n_i
}
fitting_list.append(['PSO', kwargs_pso])
kwargs_mcmc = {
'sigma_scale': 0.1,
'n_burn': 1,
'n_run': 1,
'walkerRatio': 2
}
fitting_list.append(['MCMC', kwargs_mcmc])
kwargs_mcmc['re_use_samples'] = True
fitting_list.append(['MCMC', kwargs_mcmc])
kwargs_mcmc['sampler_type'] = 'EMCEE'
fitting_list.append(['MCMC', kwargs_mcmc])
kwargs_align = {
'lowerLimit': -0.1,
'upperLimit': 0.1,
'n_particles': 2,
'n_iterations': 2
}
fitting_list.append(['align_images', kwargs_align])
kwargs_psf_iter = {
'num_iter': 2,
'psf_iter_factor': 0.5,
'stacking_method': 'mean'
}
fitting_list.append(['psf_iteration', kwargs_psf_iter])
fitting_list.append(['restart', None])
fitting_list.append(['fix_not_computed', {'free_bands': [True]}])
n_sersic_overwrite = 4
kwargs_update = {
'lens_light_add_fixed': [[0, ['n_sersic'], [n_sersic_overwrite]]],
'lens_light_remove_fixed': [[0, ['center_x']]],
'change_source_lower_limit': [[0, ['n_sersic'], [0.1]]],
'change_source_upper_limit': [[0, ['n_sersic'], [10]]]
}
fitting_list.append(['update_settings', kwargs_update])
#kwargs_model = {}, kwargs_constraints = {}, kwargs_likelihood = {}, lens_add_fixed = [],
#source_add_fixed = [], lens_light_add_fixed = [], ps_add_fixed = [], cosmo_add_fixed = [], lens_remove_fixed = [],
#source_remove_fixed = [], lens_light_remove_fixed = [], ps_remove_fixed = [], cosmo_remove_fixed = []
chain_list = fittingSequence.fit_sequence(fitting_list)
lens_fixed, source_fixed, lens_light_fixed, ps_fixed, special_fixed, extinction_fixed = fittingSequence._updateManager._fixed_kwargs
kwargs_result = fittingSequence.best_fit(bijective=False)
npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'],
self.kwargs_lens[0]['theta_E'],
decimal=1)
npt.assert_almost_equal(
fittingSequence._updateManager._lens_light_fixed[0]['n_sersic'],
n_sersic_overwrite,
decimal=8)
npt.assert_almost_equal(lens_light_fixed[0]['n_sersic'], 4, decimal=-1)
assert fittingSequence._updateManager._lower_kwargs[1][0][
'n_sersic'] == 0.1
assert fittingSequence._updateManager._upper_kwargs[1][0][
'n_sersic'] == 10
# Nested sampler tests
# further decrease the parameter space for nested samplers to run faster
fitting_list2 = []
kwargs_update2 = {
'ps_add_fixed': [[0, ['ra_source', 'dec_source'], [0, 0]]],
'lens_light_add_fixed': [[
0, ['n_sersic', 'R_sersic', 'center_x', 'center_y'],
[4, .1, 0, 0]
]],
'source_add_fixed': [[
0, ['R_sersic', 'e1', 'e2', 'center_x', 'center_y'],
[.6, .1, .1, 0, 0]
]],
'lens_add_fixed': [[
0, ['gamma', 'theta_E', 'e1', 'e2', 'center_x', 'center_y'],
[1.8, 1., .1, .1, 0, 0]
], [1, ['e1', 'e2'], [0.01, 0.01]]],
'change_source_lower_limit': [[0, ['n_sersic'], [2.9]]],
'change_source_upper_limit': [[0, ['n_sersic'], [3.1]]]
}
fitting_list2.append(['update_settings', kwargs_update2])
kwargs_multinest = {
'sampler_type': 'MULTINEST',
'kwargs_run': {
'n_live_points': 10,
'evidence_tolerance': 0.5,
'sampling_efficiency':
0.8, # 1 for posterior-only, 0 for evidence-only
'importance_nested_sampling': False,
'multimodal': True,
'const_efficiency_mode':
False, # reduce sampling_efficiency to 5% when True
},
'remove_output_dir': True,
}
fitting_list2.append(['nested_sampling', kwargs_multinest])
kwargs_dynesty = {
'sampler_type': 'DYNESTY',
'kwargs_run': {
'dlogz_init': 0.01,
'nlive_init': 3,
'nlive_batch': 3,
'maxbatch': 1,
},
}
fitting_list2.append(['nested_sampling', kwargs_dynesty])
kwargs_dypolychord = {
'sampler_type': 'DYPOLYCHORD',
'kwargs_run': {
'ninit': 8,
'nlive_const': 10,
#'seed_increment': 1,
'resume_dyn_run': False,
#'init_step': 10,
},
'polychord_settings': {
'seed': 1,
#'num_repeats': 20
},
'dypolychord_dynamic_goal':
0.8, # 1 for posterior-only, 0 for evidence-only
'remove_output_dir': True,
}
fitting_list2.append(['nested_sampling', kwargs_dypolychord])
chain_list2 = fittingSequence.fit_sequence(fitting_list2)
kwargs_fixed = fittingSequence._updateManager._fixed_kwargs
npt.assert_almost_equal(kwargs_fixed[0][1]['e1'], 0.01, decimal=2)
assert fittingSequence._updateManager._lower_kwargs[1][0][
'n_sersic'] == 2.9
assert fittingSequence._updateManager._upper_kwargs[1][0][
'n_sersic'] == 3.1
kwargs_test = {'kwargs_lens': 1}
fittingSequence.update_state(kwargs_test)
kwargs_out = fittingSequence.best_fit(bijective=True)
assert kwargs_out['kwargs_lens'] == 1
def test_minimizer(self):
n_p = 2
n_i = 2
fitting_list = []
kwargs_simplex = {'n_iterations': n_i, 'method': 'Nelder-Mead'}
fitting_list.append(['SIMPLEX', kwargs_simplex])
kwargs_simplex = {'n_iterations': n_i, 'method': 'Powell'}
fitting_list.append(['SIMPLEX', kwargs_simplex])
kwargs_pso = {
'sigma_scale': 1,
'n_particles': n_p,
'n_iterations': n_i
}
fitting_list.append(['PSO', kwargs_pso])
kwargs_mcmc = {
'sigma_scale': 1,
'n_burn': 1,
'n_run': 1,
'n_walkers': 10,
'sampler_type': 'EMCEE'
}
fitting_list.append(['MCMC', kwargs_mcmc])
kwargs_mcmc['re_use_samples'] = True
kwargs_mcmc['init_samples'] = np.array([[np.random.normal(1, 0.001)]
for i in range(100)])
fitting_list.append(['MCMC', kwargs_mcmc])
def custom_likelihood(kwargs_lens,
kwargs_source=None,
kwargs_lens_light=None,
kwargs_ps=None,
kwargs_special=None,
kwargs_extinction=None):
theta_E = kwargs_lens[0]['theta_E']
return -(theta_E - 1.)**2 / 0.1**2 / 2
kwargs_likelihood = {'custom_logL_addition': custom_likelihood}
kwargs_data_joint = {'multi_band_list': []}
kwargs_model = {'lens_model_list': ['SIS']}
kwargs_constraints = {}
lens_param = [{
'theta_E': 1,
'center_x': 0,
'center_y': 0
}], [{
'theta_E': 0.1,
'center_x': 0.1,
'center_y': 0.1
}], [{
'center_x': 0,
'center_y': 0
}], [{
'theta_E': 0,
'center_x': -10,
'center_y': -10
}], [{
'theta_E': 10,
'center_x': 10,
'center_y': 10
}]
kwargs_params = {'lens_model': lens_param}
fittingSequence = FittingSequence(kwargs_data_joint, kwargs_model,
kwargs_constraints,
kwargs_likelihood, kwargs_params)
args = fittingSequence.param_class.kwargs2args(
kwargs_lens=[{
'theta_E': 1,
'center_x': 0,
'center_y': 0
}])
kwargs_result = fittingSequence.param_class.args2kwargs(args)
print(kwargs_result)
print(args, 'test args')
chain_list = fittingSequence.fit_sequence(fitting_list)
kwargs_result = fittingSequence.best_fit(bijective=False)
npt.assert_almost_equal(kwargs_result['kwargs_lens'][0]['theta_E'],
1,
decimal=2)
def main():
args = script_utils.parse_inference_args()
test_cfg = TestConfig.from_file(args.test_config_file_path)
baobab_cfg = BaobabConfig.from_file(test_cfg.data.test_baobab_cfg_path)
cfg = TrainValConfig.from_file(test_cfg.train_val_config_file_path)
# Set device and default data type
device = torch.device(test_cfg.device_type)
if device.type == 'cuda':
torch.set_default_tensor_type('torch.cuda.' + cfg.data.float_type)
else:
torch.set_default_tensor_type('torch.' + cfg.data.float_type)
script_utils.seed_everything(test_cfg.global_seed)
############
# Data I/O #
############
train_data = XYData(
is_train=True,
Y_cols=cfg.data.Y_cols,
float_type=cfg.data.float_type,
define_src_pos_wrt_lens=cfg.data.define_src_pos_wrt_lens,
rescale_pixels=cfg.data.rescale_pixels,
rescale_pixels_type=cfg.data.rescale_pixels_type,
log_pixels=cfg.data.log_pixels,
add_pixel_noise=cfg.data.add_pixel_noise,
eff_exposure_time=cfg.data.eff_exposure_time,
train_Y_mean=None,
train_Y_std=None,
train_baobab_cfg_path=cfg.data.train_baobab_cfg_path,
val_baobab_cfg_path=None,
for_cosmology=False)
# Define val data and loader
test_data = XYData(
is_train=False,
Y_cols=cfg.data.Y_cols,
float_type=cfg.data.float_type,
define_src_pos_wrt_lens=cfg.data.define_src_pos_wrt_lens,
rescale_pixels=cfg.data.rescale_pixels,
rescale_pixels_type=cfg.data.rescale_pixels_type,
log_pixels=cfg.data.log_pixels,
add_pixel_noise=cfg.data.add_pixel_noise,
eff_exposure_time=cfg.data.eff_exposure_time,
train_Y_mean=train_data.train_Y_mean,
train_Y_std=train_data.train_Y_std,
train_baobab_cfg_path=cfg.data.train_baobab_cfg_path,
val_baobab_cfg_path=test_cfg.data.test_baobab_cfg_path,
for_cosmology=True)
master_truth = test_data.Y_df
master_truth = metadata_utils.add_qphi_columns(master_truth)
master_truth = metadata_utils.add_gamma_psi_ext_columns(master_truth)
# Figure out how many lenses BNN will predict on (must be consecutive)
if test_cfg.data.lens_indices is None:
if args.lens_indices_path is None:
# Test on all n_test lenses in the test set
n_test = test_cfg.data.n_test
lens_range = range(n_test)
else:
# Test on the lens indices in a text file at the specified path
lens_range = []
with open(args.lens_indices_path, "r") as f:
for line in f:
lens_range.append(int(line.strip()))
n_test = len(lens_range)
print("Performing H0 inference on {:d} specified lenses...".format(
n_test))
else:
if args.lens_indices_path is None:
# Test on the lens indices specified in the test config file
lens_range = test_cfg.data.lens_indices
n_test = len(lens_range)
print("Performing H0 inference on {:d} specified lenses...".format(
n_test))
else:
raise ValueError(
"Specific lens indices were specified in both the test config file and the command-line argument."
)
batch_size = max(lens_range) + 1
test_loader = DataLoader(test_data,
batch_size=batch_size,
shuffle=False,
drop_last=True)
# Output directory into which the H0 histograms and H0 samples will be saved
out_dir = test_cfg.out_dir
if not os.path.exists(out_dir):
os.makedirs(out_dir)
print("Destination folder path: {:s}".format(out_dir))
else:
raise OSError("Destination folder already exists.")
#####################
# Parameter penalty #
#####################
# Instantiate original loss function with all BNN-predicted params
orig_Y_cols = cfg.data.Y_cols
loss_fn = getattr(h0rton.losses,
cfg.model.likelihood_class)(Y_dim=test_data.Y_dim,
device=device)
# Not all predicted params will be sampled via MCMC
params_to_remove = [] #'lens_light_R_sersic', 'src_light_R_sersic']
mcmc_Y_cols = [col for col in orig_Y_cols if col not in params_to_remove]
mcmc_Y_dim = len(mcmc_Y_cols)
# Instantiate loss function with just the MCMC params
mcmc_loss_fn = getattr(h0rton.losses, cfg.model.likelihood_class)(
Y_dim=test_data.Y_dim - len(params_to_remove), device=device)
remove_param_idx, remove_idx = mcmc_utils.get_idx_for_params(
mcmc_loss_fn.out_dim, orig_Y_cols, params_to_remove,
cfg.model.likelihood_class)
mcmc_train_Y_mean = np.delete(train_data.train_Y_mean, remove_param_idx)
mcmc_train_Y_std = np.delete(train_data.train_Y_std, remove_param_idx)
parameter_penalty = mcmc_utils.HybridBNNPenalty(
mcmc_Y_cols, cfg.model.likelihood_class, mcmc_train_Y_mean,
mcmc_train_Y_std, test_cfg.h0_posterior.exclude_velocity_dispersion,
device)
custom_logL_addition = parameter_penalty.evaluate
null_spread = False
###################
# BNN predictions #
###################
# Instantiate BNN model
net = getattr(h0rton.models,
cfg.model.architecture)(num_classes=loss_fn.out_dim,
dropout_rate=cfg.model.dropout_rate)
net.to(device)
# Load trained weights from saved state
net, epoch = train_utils.load_state_dict_test(test_cfg.state_dict_path,
net, cfg.optim.n_epochs,
device)
# When only generating BNN predictions (and not running MCMC), we can afford more n_dropout
# otherwise, we fix n_dropout = mcmc_Y_dim + 1
if test_cfg.export.pred:
n_dropout = 20
n_samples_per_dropout = test_cfg.numerics.mcmc.walkerRatio
else:
n_walkers = test_cfg.numerics.mcmc.walkerRatio * (
mcmc_Y_dim + 1) # (BNN params + D_dt) times walker ratio
n_dropout = n_walkers // test_cfg.numerics.mcmc.walkerRatio
n_samples_per_dropout = test_cfg.numerics.mcmc.walkerRatio
# Initialize arrays that will store samples and BNN predictions
init_pos = np.empty(
[batch_size, n_dropout, n_samples_per_dropout, mcmc_Y_dim])
mcmc_pred = np.empty([batch_size, n_dropout, mcmc_loss_fn.out_dim])
with torch.no_grad():
net.train()
# Send some empty forward passes through the test data without backprop to adjust batchnorm weights
# (This is often not necessary. Beware if using for just 1 lens.)
for nograd_pass in range(5):
for X_, Y_ in test_loader:
X = X_.to(device)
_ = net(X)
# Obtain MC dropout samples
for d in range(n_dropout):
net.eval()
for X_, Y_ in test_loader:
X = X_.to(device)
Y = Y_.to(device)
pred = net(X)
break
mcmc_pred_d = pred.cpu().numpy()
# Replace BNN posterior's primary gaussian mean with truth values
if test_cfg.lens_posterior_type == 'default_with_truth_mean':
mcmc_pred_d[:, :len(mcmc_Y_cols)] = Y[:, :len(mcmc_Y_cols
)].cpu().numpy()
# Leave only the MCMC parameters in pred
mcmc_pred_d = mcmc_utils.remove_parameters_from_pred(
mcmc_pred_d, remove_idx, return_as_tensor=False)
# Populate pred that will define the MCMC penalty function
mcmc_pred[:, d, :] = mcmc_pred_d
# Instantiate posterior to generate BNN samples, which will serve as initial positions for walkers
bnn_post = getattr(h0rton.h0_inference.gaussian_bnn_posterior_cpu,
loss_fn.posterior_name + 'CPU')(
mcmc_Y_dim, mcmc_train_Y_mean,
mcmc_train_Y_std)
bnn_post.set_sliced_pred(mcmc_pred_d)
init_pos[:, d, :, :] = bnn_post.sample(
n_samples_per_dropout, sample_seed=test_cfg.global_seed +
d) # contains just the lens model params, no D_dt
gc.collect()
# Terminate right after generating BNN predictions (no MCMC)
if test_cfg.export.pred:
import sys
samples_path = os.path.join(out_dir, 'samples.npy')
np.save(samples_path, init_pos)
sys.exit()
#############
# MCMC loop #
#############
# Convolve MC dropout iterates with aleatoric samples
init_pos = init_pos.transpose(0, 3, 1, 2).reshape(
[batch_size, mcmc_Y_dim,
-1]).transpose(0, 2, 1) # [batch_size, n_samples, mcmc_Y_dim]
init_D_dt = np.random.uniform(0.0,
15000.0,
size=(batch_size, n_walkers, 1))
pred_mean = np.mean(init_pos, axis=1) # [batch_size, mcmc_Y_dim]
# Define assumed model profiles
kwargs_model = dict(lens_model_list=['PEMD', 'SHEAR'],
point_source_model_list=['SOURCE_POSITION'],
source_light_model_list=['SERSIC_ELLIPSE'])
astro_sig = test_cfg.image_position_likelihood.sigma # astrometric uncertainty
# Get H0 samples for each system
if not test_cfg.time_delay_likelihood.baobab_time_delays:
if 'abcd_ordering_i' not in master_truth:
raise ValueError(
"If the time delay measurements were not generated using Baobab, the user must specify the order of image positions in which the time delays are listed, in order of increasing dec."
)
kwargs_lens_eqn_solver = {
'min_distance':
0.05,
'search_window':
baobab_cfg.instrument['pixel_scale'] * baobab_cfg.image['num_pix'],
'num_iter_max':
200
}
total_progress = tqdm(total=n_test)
realized_time_delays = pd.read_csv(
test_cfg.error_model.realized_time_delays, index_col=None)
# For each lens system...
for i, lens_i in enumerate(lens_range):
# Each lens gets a unique random state for time delay measurement error realizations.
#rs_lens = np.random.RandomState(lens_i) # replaced with externally rendered time delays
###########################
# Relevant data and prior #
###########################
data_i = master_truth.iloc[lens_i].copy()
# Set BNN pred defining parameter penalty for this lens, batch processes across n_dropout
parameter_penalty.set_bnn_post_params(mcmc_pred[lens_i, :, :])
# Initialize lens model params walkers at the predictive mean
init_info = dict(
zip(mcmc_Y_cols,
pred_mean[lens_i, :] * mcmc_train_Y_std + mcmc_train_Y_mean))
lcdm = LCDM(z_lens=data_i['z_lens'],
z_source=data_i['z_src'],
flat=True)
true_img_dec = literal_eval(data_i['y_image'])
n_img = len(true_img_dec)
measured_td_sig = test_cfg.time_delay_likelihood.sigma
measured_td_wrt0 = np.array(
literal_eval(
realized_time_delays.iloc[lens_i]['measured_td_wrt0']))
kwargs_data_joint = dict(
time_delays_measured=measured_td_wrt0,
time_delays_uncertainties=measured_td_sig,
)
#############################
# Parameter init and bounds #
#############################
lens_kwargs = mcmc_utils.get_lens_kwargs(init_info,
null_spread=null_spread)
ps_kwargs = mcmc_utils.get_ps_kwargs_src_plane(init_info, astro_sig)
src_light_kwargs = mcmc_utils.get_light_kwargs(
init_info['src_light_R_sersic'], null_spread=null_spread)
special_kwargs = mcmc_utils.get_special_kwargs(
n_img, astro_sig
) # image position offset and time delay distance, aka the "special" parameters
kwargs_params = {
'lens_model': lens_kwargs,
'point_source_model': ps_kwargs,
'source_model': src_light_kwargs,
'special': special_kwargs,
}
if test_cfg.numerics.solver_type == 'NONE':
solver_type = 'NONE'
else:
solver_type = 'PROFILE_SHEAR' if n_img == 4 else 'CENTER'
#solver_type = 'NONE'
kwargs_constraints = {
'num_point_source_list': [n_img],
'Ddt_sampling': True,
'solver_type': solver_type,
}
kwargs_likelihood = {
'time_delay_likelihood': True,
'sort_images_by_dec': True,
'prior_lens': [],
'prior_special': [],
'check_bounds': True,
'check_matched_source_position': False,
'source_position_tolerance': 0.01,
'source_position_sigma': 0.01,
'source_position_likelihood': False,
'custom_logL_addition': custom_logL_addition,
'kwargs_lens_eqn_solver': kwargs_lens_eqn_solver
}
###########################
# MCMC posterior sampling #
###########################
fitting_seq = FittingSequence(kwargs_data_joint,
kwargs_model,
kwargs_constraints,
kwargs_likelihood,
kwargs_params,
verbose=False,
mpi=False)
if i == 0:
param_class = fitting_seq._updateManager.param_class
n_params, param_class_Y_cols = param_class.num_param()
init_pos = mcmc_utils.reorder_to_param_class(
mcmc_Y_cols, param_class_Y_cols, init_pos, init_D_dt)
# MCMC sample from the post-processed BNN posterior jointly with cosmology
lens_i_start_time = time.time()
if test_cfg.lens_posterior_type == 'default':
test_cfg.numerics.mcmc.update(init_samples=init_pos[lens_i, :, :])
fitting_kwargs_list_mcmc = [['MCMC', test_cfg.numerics.mcmc]]
#try:
with script_utils.HiddenPrints():
chain_list_mcmc = fitting_seq.fit_sequence(
fitting_kwargs_list_mcmc)
kwargs_result_mcmc = fitting_seq.best_fit()
lens_i_end_time = time.time()
inference_time = (lens_i_end_time - lens_i_start_time) / 60.0 # min
#############################
# Plotting the MCMC samples #
#############################
# sampler_type : 'EMCEE'
# samples_mcmc : np.array of shape `[n_mcmc_eval, n_params]`
# param_mcmc : list of str of length n_params, the parameter names
sampler_type, samples_mcmc, param_mcmc, _ = chain_list_mcmc[0]
new_samples_mcmc = mcmc_utils.postprocess_mcmc_chain(
kwargs_result_mcmc, samples_mcmc, kwargs_model, lens_kwargs[2],
ps_kwargs[2], src_light_kwargs[2], special_kwargs[2],
kwargs_constraints)
# Plot D_dt histogram
D_dt_samples = new_samples_mcmc['D_dt'].values
true_D_dt = lcdm.D_dt(H_0=data_i['H0'], Om0=0.3)
data_i['D_dt'] = true_D_dt
# Export D_dt samples for this lens
lens_inference_dict = dict(
D_dt_samples=D_dt_samples, # kappa_ext=0 for these samples
inference_time=inference_time,
true_D_dt=true_D_dt,
)
lens_inference_dict_save_path = os.path.join(
out_dir, 'D_dt_dict_{0:04d}.npy'.format(lens_i))
np.save(lens_inference_dict_save_path, lens_inference_dict)
# Optionally export the MCMC samples
if test_cfg.export.mcmc_samples:
mcmc_samples_path = os.path.join(
out_dir, 'mcmc_samples_{0:04d}.csv'.format(lens_i))
new_samples_mcmc.to_csv(mcmc_samples_path, index=None)
# Optionally export the D_dt histogram
if test_cfg.export.D_dt_histogram:
cleaned_D_dt_samples = h0_utils.remove_outliers_from_lognormal(
D_dt_samples, 3)
_ = plotting_utils.plot_D_dt_histogram(cleaned_D_dt_samples,
lens_i,
true_D_dt,
save_dir=out_dir)
# Optionally export the plot of MCMC chain
if test_cfg.export.mcmc_chain:
mcmc_chain_path = os.path.join(
out_dir, 'mcmc_chain_{0:04d}.png'.format(lens_i))
plotting_utils.plot_mcmc_chain(chain_list_mcmc, mcmc_chain_path)
# Optionally export posterior cornerplot of select lens model parameters with D_dt
if test_cfg.export.mcmc_corner:
mcmc_corner_path = os.path.join(
out_dir, 'mcmc_corner_{0:04d}.png'.format(lens_i))
plotting_utils.plot_mcmc_corner(
new_samples_mcmc[test_cfg.export.mcmc_cols],
data_i[test_cfg.export.mcmc_cols],
test_cfg.export.mcmc_col_labels, mcmc_corner_path)
total_progress.update(1)
gc.collect()
realized_time_delays.to_csv(os.path.join(out_dir,
'realized_time_delays.csv'),
index=None)
total_progress.close()
def swim(self, lens_name, model_id, log=True, mpi=False,
recipe_name='default', sampler='EMCEE', thread_count=1):
"""
Run models for a single lens.
:param lens_name: lens name
:type lens_name: `str`
:param model_id: identifier for the model run
:type model_id: `str`
:param log: if `True`, all `print` statements will be logged
:type log: `bool`
:param mpi: MPI option
:type mpi: `bool`
:param recipe_name: recipe for pre-sampling optimization, supported
ones now: 'default' and 'galaxy-galaxy'
:type recipe_name: `str`
:param sampler: 'EMCEE' or 'COSMOHAMMER', cosmohammer is kept for
legacy
:type sampler: `str`
:param thread_count: number of threads if `multiprocess` is used
:type thread_count: `int`
:return:
:rtype:
"""
pool = choose_pool(mpi=mpi)
if log and pool.is_master():
log_file = open(self.file_system.get_log_file_path(lens_name,
model_id), 'wt')
sys.stdout = log_file
config = self.get_lens_config(lens_name)
recipe = Recipe(config, sampler=sampler, thread_count=thread_count)
psf_supersampling_factor = config.get_psf_supersampled_factor()
kwargs_data_joint = self.get_kwargs_data_joint(
lens_name,
psf_supersampled_factor=psf_supersampling_factor)
fitting_sequence = FittingSequence(
kwargs_data_joint,
config.get_kwargs_model(),
config.get_kwargs_constraints(),
config.get_kwargs_likelihood(),
config.get_kwargs_params(),
mpi=mpi
)
fitting_kwargs_list = recipe.get_recipe(
kwargs_data_joint=kwargs_data_joint,
recipe_name=recipe_name)
fit_output = fitting_sequence.fit_sequence(fitting_kwargs_list)
kwargs_result = fitting_sequence.best_fit(bijective=False)
output = {
'settings': config.settings,
'kwargs_result': kwargs_result,
'fit_output': fit_output,
}
if pool.is_master():
self.file_system.save_output(lens_name, model_id, output)
if log and pool.is_master():
log_file.close()
def main():
args = parse_args()
test_cfg = TestConfig.from_file(args.test_config_file_path)
train_val_cfg = TrainValConfig.from_file(
test_cfg.train_val_config_file_path)
# Set device and default data type
device = torch.device(test_cfg.device_type)
if device.type == 'cuda':
torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
seed_everything(test_cfg.global_seed)
############
# Data I/O #
############
test_data = TDLMCData(data_cfg=train_val_cfg.data, rung_i=args.rung_idx)
master_truth = test_data.cosmo_df
if test_cfg.data.lens_indices is None:
if args.lens_indices_path is None:
# Test on all n_test lenses in the test set
n_test = test_cfg.data.n_test
lens_range = range(n_test)
else:
# Test on the lens indices in a text file at the specified path
lens_range = []
with open(args.lens_indices_path, "r") as f:
for line in f:
lens_range.append(int(line.strip()))
n_test = len(lens_range)
print("Performing H0 inference on {:d} specified lenses...".format(
n_test))
else:
if args.lens_indices_path is None:
# Test on the lens indices specified in the test config file
lens_range = test_cfg.data.lens_indices
n_test = len(lens_range)
print("Performing H0 inference on {:d} specified lenses...".format(
n_test))
else:
raise ValueError(
"Specific lens indices were specified in both the test config file and the command-line argument."
)
batch_size = max(lens_range) + 1
test_loader = DataLoader(test_data,
batch_size=batch_size,
shuffle=False,
drop_last=True)
# Output directory into which the H0 histograms and H0 samples will be saved
out_dir = test_cfg.out_dir
if not os.path.exists(out_dir):
os.makedirs(out_dir)
print("Destination folder path: {:s}".format(out_dir))
else:
raise OSError("Destination folder already exists.")
######################
# Load trained state #
######################
# Instantiate loss function, to append to the MCMC objective as the prior
orig_Y_cols = train_val_cfg.data.Y_cols
loss_fn = getattr(h0rton.losses, train_val_cfg.model.likelihood_class)(
Y_dim=train_val_cfg.data.Y_dim, device=device)
# Instantiate MCMC parameter penalty function
params_to_remove = ['lens_light_R_sersic'] #, 'src_light_R_sersic']
mcmc_Y_cols = [col for col in orig_Y_cols if col not in params_to_remove]
mcmc_Y_dim = len(mcmc_Y_cols)
mcmc_loss_fn = getattr(
h0rton.losses, train_val_cfg.model.likelihood_class)(
Y_dim=train_val_cfg.data.Y_dim - len(params_to_remove),
device=device)
remove_param_idx, remove_idx = mcmc_utils.get_idx_for_params(
mcmc_loss_fn.out_dim, orig_Y_cols, params_to_remove,
train_val_cfg.model.likelihood_class)
mcmc_train_Y_mean = np.delete(train_val_cfg.data.train_Y_mean,
remove_param_idx)
mcmc_train_Y_std = np.delete(train_val_cfg.data.train_Y_std,
remove_param_idx)
parameter_penalty = mcmc_utils.HybridBNNPenalty(
mcmc_Y_cols, train_val_cfg.model.likelihood_class, mcmc_train_Y_mean,
mcmc_train_Y_std, test_cfg.h0_posterior.exclude_velocity_dispersion,
device)
custom_logL_addition = parameter_penalty.evaluate if test_cfg.lens_posterior_type.startswith(
'default') else None
null_spread = True if test_cfg.lens_posterior_type == 'truth' else False
# Instantiate model
net = getattr(
h0rton.models,
train_val_cfg.model.architecture)(num_classes=loss_fn.out_dim)
net.to(device)
# Load trained weights from saved state
net, epoch = train_utils.load_state_dict_test(test_cfg.state_dict_path,
net,
train_val_cfg.optim.n_epochs,
device)
with torch.no_grad():
net.eval()
for X_ in test_loader:
X = X_.to(device)
pred = net(X)
break
mcmc_pred = pred.cpu().numpy()
mcmc_pred = mcmc_utils.remove_parameters_from_pred(mcmc_pred,
remove_idx,
return_as_tensor=False)
# Instantiate posterior for BNN samples, to initialize the walkers
bnn_post = getattr(h0rton.h0_inference.gaussian_bnn_posterior,
loss_fn.posterior_name)(mcmc_Y_dim, device,
mcmc_train_Y_mean,
mcmc_train_Y_std)
bnn_post.set_sliced_pred(torch.tensor(mcmc_pred))
n_walkers = test_cfg.numerics.mcmc.walkerRatio * (
mcmc_Y_dim + 1) # BNN params + H0 times walker ratio
init_pos = bnn_post.sample(
n_walkers, sample_seed=test_cfg.global_seed
) # [batch_size, n_walkers, mcmc_Y_dim] contains just the lens model params, no D_dt
init_D_dt = np.random.uniform(0.0,
10000.0,
size=(batch_size, n_walkers,
1)) # FIXME: init H0 hardcoded
kwargs_model = dict(lens_model_list=['PEMD', 'SHEAR'],
point_source_model_list=['SOURCE_POSITION'],
source_light_model_list=['SERSIC_ELLIPSE'])
astro_sig = test_cfg.image_position_likelihood.sigma
# Get H0 samples for each system
if not test_cfg.time_delay_likelihood.baobab_time_delays:
if 'abcd_ordering_i' not in master_truth:
raise ValueError(
"If the time delay measurements were not generated using Baobab, the user must specify the order of image positions in which the time delays are listed, in order of increasing dec."
)
lenses_skipped = [] # keeps track of lenses that skipped MCMC
total_progress = tqdm(total=n_test)
# For each lens system...
for i, lens_i in enumerate(lens_range):
# Each lens gets a unique random state for td and vd measurement error realizations.
rs_lens = np.random.RandomState(lens_i)
###########################
# Relevant data and prior #
###########################
data_i = master_truth.iloc[lens_i].copy()
parameter_penalty.set_bnn_post_params(
mcmc_pred[lens_i, :]) # set the BNN parameters
# Init values for the lens model params
if test_cfg.lens_posterior_type == 'default':
init_info = dict(
zip(
mcmc_Y_cols,
mcmc_pred[lens_i, :len(mcmc_Y_cols)] * mcmc_train_Y_std +
mcmc_train_Y_mean)) # mean of primary Gaussian
else: # types 'hybrid_with_truth_mean' and 'truth'
init_info = dict(zip(mcmc_Y_cols,
data_i[mcmc_Y_cols].values)) # truth params
if not test_cfg.h0_posterior.exclude_velocity_dispersion:
parameter_penalty.set_vel_disp_params()
raise NotImplementedError
lcdm = LCDM(z_lens=data_i['z_lens'],
z_source=data_i['z_src'],
flat=True)
# Data is BCD - A with a certain ABCD ordering, so inferred time delays should follow this convention.
measured_td_wrt0 = np.array(data_i['measured_td']) # [n_img - 1,]
measured_td_sig = np.array(data_i['measured_td_err']) # [n_img - 1,]
abcd_ordering_i = np.array(data_i['abcd_ordering_i'])
n_img = len(abcd_ordering_i)
kwargs_data_joint = dict(
time_delays_measured=measured_td_wrt0,
time_delays_uncertainties=measured_td_sig,
abcd_ordering_i=abcd_ordering_i,
#vel_disp_measured=measured_vd, # TODO: optionally exclude
#vel_disp_uncertainty=vel_disp_sig,
)
if not test_cfg.h0_posterior.exclude_velocity_dispersion:
measured_vd = data_i['true_vd'] * (
1.0 + rs_lens.randn() *
test_cfg.error_model.velocity_dispersion_frac_error)
kwargs_data_joint['vel_disp_measured'] = measured_vd
kwargs_data_joint[
'vel_disp_sig'] = test_cfg.velocity_dispersion_likelihood.sigma
#############################
# Parameter init and bounds #
#############################
lens_kwargs = mcmc_utils.get_lens_kwargs(init_info,
null_spread=null_spread)
ps_kwargs = mcmc_utils.get_ps_kwargs_src_plane(init_info,
astro_sig,
null_spread=null_spread)
src_light_kwargs = mcmc_utils.get_light_kwargs(
init_info['src_light_R_sersic'], null_spread=null_spread)
special_kwargs = mcmc_utils.get_special_kwargs(
n_img, astro_sig, null_spread=null_spread
) # image position offset and time delay distance, aka the "special" parameters
kwargs_params = {
'lens_model': lens_kwargs,
'point_source_model': ps_kwargs,
'source_model': src_light_kwargs,
'special': special_kwargs,
}
if test_cfg.numerics.solver_type == 'NONE':
solver_type = 'NONE'
else:
solver_type = 'PROFILE_SHEAR' if n_img == 4 else 'CENTER'
#solver_type = 'NONE'
kwargs_constraints = {
'num_point_source_list': [n_img],
'Ddt_sampling': True,
'solver_type': solver_type,
}
kwargs_likelihood = {
'time_delay_likelihood': True,
'sort_images_by_dec': True,
'prior_lens': [],
'prior_special': [],
'check_bounds': True,
'check_matched_source_position': False,
'source_position_tolerance': 0.01,
'source_position_sigma': 0.01,
'source_position_likelihood': False,
'custom_logL_addition': custom_logL_addition,
}
###########################
# MCMC posterior sampling #
###########################
fitting_seq = FittingSequence(kwargs_data_joint,
kwargs_model,
kwargs_constraints,
kwargs_likelihood,
kwargs_params,
verbose=False,
mpi=False)
if i == 0:
param_class = fitting_seq._updateManager.param_class
n_params, param_class_Y_cols = param_class.num_param()
init_pos = mcmc_utils.reorder_to_param_class(
mcmc_Y_cols, param_class_Y_cols, init_pos, init_D_dt)
# MCMC sample from the post-processed BNN posterior jointly with cosmology
lens_i_start_time = time.time()
if test_cfg.lens_posterior_type == 'default':
test_cfg.numerics.mcmc.update(init_samples=init_pos[lens_i, :, :])
fitting_kwargs_list_mcmc = [['MCMC', test_cfg.numerics.mcmc]]
#with HiddenPrints():
try:
chain_list_mcmc = fitting_seq.fit_sequence(
fitting_kwargs_list_mcmc)
kwargs_result_mcmc = fitting_seq.best_fit()
except:
print("lens {:d} skipped".format(lens_i))
total_progress.update(1)
lenses_skipped.append(lens_i)
continue
lens_i_end_time = time.time()
inference_time = (lens_i_end_time - lens_i_start_time) / 60.0 # min
#############################
# Plotting the MCMC samples #
#############################
# sampler_type : 'EMCEE'
# samples_mcmc : np.array of shape `[n_mcmc_eval, n_params]`
# param_mcmc : list of str of length n_params, the parameter names
sampler_type, samples_mcmc, param_mcmc, _ = chain_list_mcmc[0]
new_samples_mcmc = mcmc_utils.postprocess_mcmc_chain(
kwargs_result_mcmc, samples_mcmc, kwargs_model, lens_kwargs[2],
ps_kwargs[2], src_light_kwargs[2], special_kwargs[2],
kwargs_constraints)
# Plot D_dt histogram
D_dt_samples = new_samples_mcmc['D_dt'].values
true_D_dt = lcdm.D_dt(H_0=data_i['H0'], Om0=0.27)
data_i['D_dt'] = true_D_dt
# Export D_dt samples for this lens
lens_inference_dict = dict(
D_dt_samples=D_dt_samples, # kappa_ext=0 for these samples
inference_time=inference_time,
true_D_dt=true_D_dt,
)
lens_inference_dict_save_path = os.path.join(
out_dir, 'D_dt_dict_{0:04d}.npy'.format(lens_i))
np.save(lens_inference_dict_save_path, lens_inference_dict)
# Optionally export the MCMC samples
if test_cfg.export.mcmc_samples:
mcmc_samples_path = os.path.join(
out_dir, 'mcmc_samples_{0:04d}.csv'.format(lens_i))
new_samples_mcmc.to_csv(mcmc_samples_path, index=None)
# Optionally export the D_dt histogram
if test_cfg.export.D_dt_histogram:
cleaned_D_dt_samples = h0_utils.remove_outliers_from_lognormal(
D_dt_samples, 3)
_ = plotting_utils.plot_D_dt_histogram(cleaned_D_dt_samples,
lens_i,
true_D_dt,
save_dir=out_dir)
# Optionally export the plot of MCMC chain
if test_cfg.export.mcmc_chain:
mcmc_chain_path = os.path.join(
out_dir, 'mcmc_chain_{0:04d}.png'.format(lens_i))
plotting_utils.plot_mcmc_chain(chain_list_mcmc, mcmc_chain_path)
# Optionally export posterior cornerplot of select lens model parameters with D_dt
if test_cfg.export.mcmc_corner:
mcmc_corner_path = os.path.join(
out_dir, 'mcmc_corner_{0:04d}.png'.format(lens_i))
plotting_utils.plot_mcmc_corner(
new_samples_mcmc[test_cfg.export.mcmc_cols], None,
test_cfg.export.mcmc_col_labels, mcmc_corner_path)
total_progress.update(1)
total_progress.close()
'n_sersic': 6,
'center_x': (x0blob + 4) * deltaPix,
'center_y': ((60 - y0blob) + 4) * deltaPix
})
lens_light_params = [
kwargs_lens_light_init, kwargs_lens_light_sigma, fixed_lens_light,
kwargs_lower_lens_light, kwargs_upper_lens_light
]
kwargs_params = {'lens_light_model': lens_light_params}
from lenstronomy.Workflow.fitting_sequence import FittingSequence
fitting_seqI = FittingSequence(multi_band_listI, kwargs_modelI,
kwargs_constraints, kwargs_likelihood,
kwargs_params)
fitting_kwargs_list = [[
'MCMC', {
'n_burn': 100,
'n_run': 100,
'walkerRatio': 10,
'sigma_scale': .1
}
]]
chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seqI.fit_sequence(
fitting_kwargs_list)
lens_result, source_result, lens_light_resultI, ps_result, cosmo_result = fitting_seqI.best_fit(
)
def sl_sys_analysis():
# Get command line arguments
args = {}
if comm_rank == 0:
print(":Registered %d processes" % comm_size)
args["infile"] = sys.argv[1]
args["nimgs"] = sys.argv[2]
args["los"] = sys.argv[3]
args = comm.bcast(args)
# Organize devision of strong lensing systems
with open(args["infile"], "r") as myfile:
limg_data = myfile.read()
systems = json.loads(limg_data)
sys_nr_per_proc = int(len(systems) / comm_size)
start_sys = sys_nr_per_proc * comm_rank
end_sys = sys_nr_per_proc * (comm_rank + 1)
with open("../lens_catalogs_sie_only.json", "r") as myfile:
limg_data = myfile.read()
systems_prior = json.loads(limg_data)
if comm_rank == 0:
print("Each process will have %d systems" % sys_nr_per_proc)
print("That should take app. %f min." % (sys_nr_per_proc * 20))
results = {"D_dt": []}
for ii in range(len(systems))[(start_sys + 2) : end_sys]:
system = systems[ii]
system_prior = systems_prior[ii]
print("Analysing system ID: %d" % system["losID"])
# the data set is
z_lens = system_prior["zl"]
z_source = 2.0
# image positions units of arcsec
# time delays units of days
# image brightness
# amplitude (in arbitrary linear units, not magnitudes)
ximg = []
yimg = []
delay = []
image_amps = []
for jj in range(system["nimgs"]):
ximg.append(system["ximg"][jj])
yimg.append(system["yimg"][jj])
delay.append(system["delay"][jj])
image_amps.append(system["mags"][jj])
ximg = np.asarray(ximg)
yimg = np.asarray(yimg)
# 1-sigma astrometric uncertainties of the image positions
# (assuming equal precision)
astrometry_sigma = 0.004
delay = np.asarray(delay)
d_dt = delay[1:] - delay[0]
# 1-sigma uncertainties in the time-delay measurement (in units of days)
d_dt_sigma = np.ones(system["nimgs"] - 1) * 2
image_amps = np.asarray(image_amps)
image_amps_sigma = np.ones(system["nimgs"]) * 0.3
flux_ratios = image_amps[1:] - image_amps[0]
flux_ratio_errors = np.ones(system["nimgs"] - 1) * 0.1
# lens model choicers
lens_model_list = ["SPEP", "SHEAR"]
fixed_lens = []
kwargs_lens_init = []
kwargs_lens_sigma = []
kwargs_lower_lens = []
kwargs_upper_lens = []
fixed_lens.append({})
kwargs_lens_init.append(
{
"theta_E": 1.0,
"gamma": 2,
"center_x": 0,
"center_y": 0,
"e1": 0,
"e2": 0.0,
}
)
kwargs_lens_sigma.append(
{
"theta_E": 0.2,
"e1": 0.1,
"e2": 0.1,
"gamma": 0.1,
"center_x": 0.1,
"center_y": 0.1,
}
)
kwargs_lower_lens.append(
{
"theta_E": 0.01,
"e1": -0.5,
"e2": -0.5,
"gamma": 1.5,
"center_x": -10,
"center_y": -10,
}
)
kwargs_upper_lens.append(
{
"theta_E": 10,
"e1": 0.5,
"e2": 0.5,
"gamma": 2.5,
"center_x": 10,
"center_y": 10,
}
)
fixed_lens.append({"ra_0": 0, "dec_0": 0})
kwargs_lens_init.append({"e1": 0.0, "e2": 0.0})
kwargs_lens_sigma.append({"e1": 0.1, "e2": 0.1})
kwargs_lower_lens.append({"e1": -0.2, "e2": -0.2})
kwargs_upper_lens.append({"e1": 0.2, "e2": 0.2})
lens_params = [
kwargs_lens_init,
kwargs_lens_sigma,
fixed_lens,
kwargs_lower_lens,
kwargs_upper_lens,
]
point_source_list = ["LENSED_POSITION"]
fixed_ps = [
{"ra_image": ximg, "dec_image": yimg}
] # we fix the image position coordinates
kwargs_ps_init = fixed_ps
kwargs_ps_sigma = [
{
"ra_image": 0.01 * np.ones(len(ximg)),
"dec_image": 0.01 * np.ones(len(ximg)),
}
]
kwargs_lower_ps = [
{
"ra_image": -10 * np.ones(len(ximg)),
"dec_image": -10 * np.ones(len(ximg)),
}
]
kwargs_upper_ps = [
{"ra_image": 10 * np.ones(len(ximg)), "dec_image": 10 * np.ones(len(ximg))}
]
ps_params = [
kwargs_ps_init,
kwargs_ps_sigma,
fixed_ps,
kwargs_lower_ps,
kwargs_upper_ps,
]
# we let some freedome in how well the actual image positions are matching those
# given by the data (indicated as 'ra_image', 'dec_image' and held fixed while fitting)
fixed_cosmo = {}
kwargs_cosmo_init = {
"D_dt": 5000,
"delta_x_image": np.zeros_like(ximg),
"delta_y_image": np.zeros_like(ximg),
}
kwargs_cosmo_sigma = {
"D_dt": 10000,
"delta_x_image": np.ones_like(ximg) * astrometry_sigma,
"delta_y_image": np.ones_like(ximg) * astrometry_sigma,
}
kwargs_lower_cosmo = {
"D_dt": 0,
"delta_x_image": np.ones_like(ximg) * (-1),
"delta_y_image": np.ones_like(ximg) * (-1),
}
kwargs_upper_cosmo = {
"D_dt": 10000,
"delta_x_image": np.ones_like(ximg) * (1),
"delta_y_image": np.ones_like(ximg) * (1),
}
cosmo_params = [
kwargs_cosmo_init,
kwargs_cosmo_sigma,
fixed_cosmo,
kwargs_lower_cosmo,
kwargs_upper_cosmo,
]
kwargs_params = {
"lens_model": lens_params,
"point_source_model": ps_params,
"cosmography": cosmo_params,
}
# ## setup options for likelihood and parameter sampling
kwargs_constraints = {
"num_point_source_list": [4],
# any proposed lens model must satisfy the image positions
# appearing at the position of the point sources being sampeld
"solver_type": "PROFILE_SHEAR",
"cosmo_type": "D_dt", # sampling of the time-delay distance
# explicit modelling of the astrometric imperfection of
# the point source positions
"point_source_offset": True,
}
kwargs_likelihood = {
"check_bounds": True,
"point_source_likelihood": True,
"position_uncertainty": astrometry_sigma,
"check_solver": True,
"solver_tolerance": 0.001,
"time_delay_likelihood": True,
"image_likelihood": False,
# this needs to be explicitly given when not having imaging data
"flux_ratio_likelihood": True, # enables the flux ratio likelihood
}
kwargs_data_joint = {
"time_delays_measured": d_dt,
"time_delays_uncertainties": d_dt_sigma,
"flux_ratios": flux_ratios,
"flux_ratio_errors": flux_ratio_errors,
}
kwargs_model = {
"lens_model_list": lens_model_list,
"point_source_model_list": point_source_list,
}
mpi = False # MPI possible, but not supported through that notebook.
fitting_seq = FittingSequence(
kwargs_data_joint,
kwargs_model,
kwargs_constraints,
kwargs_likelihood,
kwargs_params,
)
fitting_kwargs_list = [
# ['update_settings', {'lens_add_fixed': [[0, ['gamma']]]}],
["PSO", {"sigma_scale": 1.0, "n_particles": 100, "n_iterations": 100}]
]
chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(
fitting_kwargs_list
)
lens_result, source_result, lens_light_result, ps_result, cosmo_result = (
fitting_seq.best_fit()
)
# and now we run the MCMC
fitting_kwargs_list = [
["PSO", {"sigma_scale": 0.1, "n_particles": 100, "n_iterations": 100}],
[
"MCMC",
{"n_burn": 200, "n_run": 200, "walkerRatio": 10, "sigma_scale": 0.1},
],
]
chain_list, param_list, samples_mcmc, param_mcmc, dist_mcmc = fitting_seq.fit_sequence(
fitting_kwargs_list
)
lens_result, source_result, lens_light_result, ps_result, cosmo_result = (
fitting_seq.best_fit()
)
# print("number of non-linear parameters in the MCMC process: ", len(param_mcmc))
# print("parameters in order: ", param_mcmc)
print("number of evaluations in the MCMC process: ", np.shape(samples_mcmc)[0])
param = Param(
kwargs_model,
fixed_lens,
kwargs_fixed_ps=fixed_ps,
kwargs_fixed_cosmo=fixed_cosmo,
kwargs_lens_init=lens_result,
**kwargs_constraints,
)
# the number of non-linear parameters and their names #
num_param, param_list = param.num_param()
lensAnalysis = LensAnalysis(kwargs_model)
mcmc_new_list = []
labels_new = [
r"$\gamma$",
r"$\phi_{ext}$",
r"$\gamma_{ext}$",
r"$D_{\Delta t}$",
]
for i in range(len(samples_mcmc)):
# transform the parameter position of the MCMC chain in a
# lenstronomy convention with keyword arguments #
kwargs_lens_out, kwargs_light_source_out, kwargs_light_lens_out, kwargs_ps_out, kwargs_cosmo = param.args2kwargs(
samples_mcmc[i]
)
D_dt = kwargs_cosmo["D_dt"]
gamma = kwargs_lens_out[0]["gamma"]
phi_ext, gamma_ext = lensAnalysis._lensModelExtensions.external_shear(
kwargs_lens_out
)
mcmc_new_list.append([gamma, phi_ext, gamma_ext, D_dt])
# plot = corner.corner(mcmc_new_list, labels=labels_new, show_titles=True)
# and here the predicted angular diameter distance from a
# default cosmology (attention for experimenter bias!)
lensCosmo = LensCosmo(z_lens=z_lens, z_source=z_source)
results["D_dt"].append(D_dt)
f = h5py.File(
"H0_"
+ args["los"]
+ "_nimgs"
+ str(args["nimgs"])
+ "_"
+ str(comm_rank)
+ ".hdf5",
"w",
)
hf.create_dataset("D_dt", data=results["D_dt"])
hf.close()