def test_source(self):
multi_band_list = [[
self.kwargs_data, self.kwargs_psf, self.kwargs_numerics
]]
lensPlot = ModelPlot(multi_band_list,
self.kwargs_model,
self.kwargs_params,
arrow_size=0.02,
cmap_string="gist_heat")
source, coords_source = lensPlot.source(band_index=0,
numPix=10,
deltaPix=0.1,
image_orientation=True)
assert len(source) == 10
source, coords_source = lensPlot.source(band_index=0,
numPix=10,
deltaPix=0.1,
image_orientation=False)
assert len(source) == 10
source, coords_source = lensPlot.source(band_index=0,
numPix=10,
deltaPix=0.1,
center=[0, 0])
assert len(source) == 10
def test_reconstruction_all_bands(self):
multi_band_list = [[
self.kwargs_data, self.kwargs_psf, self.kwargs_numerics
], [self.kwargs_data, self.kwargs_psf, self.kwargs_numerics]]
lensPlot = ModelPlot(multi_band_list,
self.kwargs_model,
self.kwargs_params,
arrow_size=0.02,
cmap_string="gist_heat",
multi_band_type='joint-linear',
bands_compute=[True, True])
f, axes = lensPlot.reconstruction_all_bands()
assert len(axes) == 2
assert len(axes[0]) == 3
plt.close()
def test_joint_linear(self):
multi_band_list = [[
self.kwargs_data, self.kwargs_psf, self.kwargs_numerics
], [self.kwargs_data, self.kwargs_psf, self.kwargs_numerics]]
lensPlot = ModelPlot(multi_band_list,
self.kwargs_model,
self.kwargs_params,
arrow_size=0.02,
cmap_string="gist_heat",
multi_band_type='joint-linear',
bands_compute=[True, False])
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.data_plot(ax=ax,
numPix=10,
deltaPix_source=0.1,
v_min=None,
v_max=None,
with_caustics=False,
caustic_color='yellow',
fsize=15,
plot_scale='linear')
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.model_plot(ax=ax,
numPix=10,
deltaPix_source=0.1,
v_min=None,
v_max=None,
with_caustics=False,
caustic_color='yellow',
fsize=15,
plot_scale='linear')
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.convergence_plot(ax=ax,
numPix=10,
deltaPix_source=0.1,
v_min=None,
v_max=None,
with_caustics=False,
caustic_color='yellow',
fsize=15,
plot_scale='linear')
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.normalized_residual_plot(ax=ax)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.magnification_plot(ax=ax)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.decomposition_plot(ax=ax)
plt.close()
def test_source_plot(self):
multi_band_list = [[
self.kwargs_data, self.kwargs_psf, self.kwargs_numerics
]]
lensPlot = ModelPlot(multi_band_list,
self.kwargs_model,
self.kwargs_params,
arrow_size=0.02,
cmap_string="gist_heat")
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.source_plot(ax=ax,
numPix=10,
deltaPix_source=0.1,
v_min=None,
v_max=None,
with_caustics=False,
caustic_color='yellow',
fsize=15,
plot_scale='linear')
plt.close()
def get_model_plot(self,
lens_name,
model_id=None,
kwargs_result=None,
band_index=0,
data_cmap='cubehelix'):
"""
Get the `ModelPlot` instance from lenstronomy for the lens.
:param lens_name: name of the lens
:type lens_name: `str`
:param model_id: model run identifier
:type model_id: `str`
:param kwargs_result: lenstronomy `kwargs_result` dictionary. If
provided, it will be used to plot the model, otherwise the model
will be plotted from the saved/loaded outputs for `lens_name` and
`model_id`.
:type kwargs_result: `dict`
:param band_index: index of band to plot for multi-band case
:type band_index: `int`
:param data_cmap: colormap for image, reconstruction, and source plots
:type data_cmap: `str` or `matplotlib.colors.Colormap`
:return: `ModelPlot` instance, maximum pixel value of the image
:rtype: `obj`, `float`
"""
if model_id is None and kwargs_result is None:
raise ValueError('Either the `model_id` or the `kwargs_result` '
'needs to be provided!')
if kwargs_result is None:
self.load_output(lens_name, model_id)
kwargs_result = self.kwargs_result
multi_band_list_out = self.get_kwargs_data_joint(
lens_name)['multi_band_list']
config = ModelConfig(settings=self.model_settings)
mask = config.get_masks()
kwargs_model = config.get_kwargs_model()
v_max = np.log10(
multi_band_list_out[band_index][0]['image_data'].max())
model_plot = ModelPlot(multi_band_list_out,
kwargs_model,
kwargs_result,
arrow_size=0.02,
cmap_string=data_cmap,
likelihood_mask_list=mask,
multi_band_type='multi-linear')
return model_plot, v_max
def test_lensModelPlot(self):
multi_band_list = [[
self.kwargs_data, self.kwargs_psf, self.kwargs_numerics
]]
lensPlot = ModelPlot(multi_band_list,
self.kwargs_model,
self.kwargs_params,
arrow_size=0.02,
cmap_string="gist_heat",
multi_band_type='single-band')
lensPlot.plot_main(with_caustics=True)
plt.close()
cmap = plt.get_cmap('gist_heat')
lensPlot = ModelPlot(multi_band_list,
self.kwargs_model,
self.kwargs_params,
arrow_size=0.02,
cmap_string=cmap)
lensPlot.plot_separate()
plt.close()
lensPlot.plot_subtract_from_data_all()
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
lensPlot.deflection_plot(ax=ax, with_caustics=True, axis=1)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
lensPlot.subtract_from_data_plot(ax=ax)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
lensPlot.deflection_plot(ax=ax, with_caustics=True, axis=0)
plt.close()
numPix = 100
deltaPix_source = 0.01
f, ax = plt.subplots(1, 1, figsize=(4, 4))
lensPlot.error_map_source_plot(ax=ax,
numPix=numPix,
deltaPix_source=deltaPix_source,
with_caustics=True)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
lensPlot.absolute_residual_plot(ax=ax)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
lensPlot.plot_extinction_map(ax=ax)
plt.close()
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
)
fix_setting = [fixed_lens, fixed_source, fixed_lens_light, None, None]
# sampler_type, samples_mcmc, param_mcmc, dist_mcmc = chain_list[-1]
mcmc_new_list = []
pickle.dump([
multi_band_list, kwargs_model, kwargs_result, chain_list,
fix_setting, mcmc_new_list
], open(folder + savename, 'wb'))
#Print fitting result:
multi_band_list, kwargs_model, kwargs_result, chain_list, fix_setting, _ = pickle.load(
open(folder + savename, 'rb'))
fixed_lens, fixed_source, fixed_lens_light, fixed_ps, fixed_cosmo = fix_setting
labels_new = [r"$\gamma$", r"$e1$", r"$e2$"]
modelPlot = ModelPlot(multi_band_list,
kwargs_model,
kwargs_result,
arrow_size=0.02,
cmap_string="gist_heat")
f, axes = modelPlot.plot_main()
f.show()
f, axes = modelPlot.plot_separate()
f.show()
f, axes = modelPlot.plot_subtract_from_data_all()
f.show()
sampler_type, samples_mcmc, param_mcmc, dist_mcmc = chain_list[-1]
# for i in range(len(chain_list)):
# chain_plot.plot_chain_list(chain_list, i)
param = Param(kwargs_model,
fixed_lens,
def model_plot(self, save_plot=False, show_plot=True):
"""
Show the fitting plot based on lenstronomy.Plots.model_plot.ModelPlot
"""
# this is the linear inversion. The kwargs will be updated afterwards
modelPlot = ModelPlot(
self.fitting_specify_class.kwargs_data_joint['multi_band_list'],
self.fitting_specify_class.kwargs_model,
self.kwargs_result,
arrow_size=0.02,
cmap_string="gist_heat",
likelihood_mask_list=self.fitting_specify_class.
kwargs_likelihood['image_likelihood_mask_list'])
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',
lens_light_add=True,
unconvolved=True)
modelPlot.decomposition_plot(ax=axes[1, 1],
text='Host galaxy convolved',
lens_light_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',
lens_light_add=True)
modelPlot.subtract_from_data_plot(
ax=axes[2, 2],
text='Data - host galaxy - Point Source',
lens_light_add=True,
point_source_add=True)
f.tight_layout()
if save_plot == True:
plt.savefig('{0}_model.pdf'.format(self.savename))
if show_plot == True:
plt.show()
else:
plt.close()
def plot_modeling(self,
kwargs_result,
center=[0, 0],
deltaPix_s=0.03,
numPix_s=None,
text_source='',
data_index=0,
text='sys',
img_name='sys',
font_size=25,
scale_size=0.1,
fig_close=False,
likelihood_mask_list=None):
"""
a function to show modeling process containing data, reconstructed image, residual map,
and reconstructed source.
:param kwargs_result: modeling results
:param deltaPix: pixel scale in the source plane
:param numPix: pixel numbers in the source plane
:param multi_band_type:string, e.g., 'joint-linear', 'single-band'
:param text: string, label of reconstructed image
:param text_source:string, label of reconstructed source
:param img_name: string, label of saved images
:param font_size: font_size
:return:
"""
model_plot = ModelPlot(self.multi_band_list,
self.kwargs_model,
kwargs_result,
arrow_size=0.02,
cmap_string="gist_heat",
multi_band_type=self.multi_band_type,
likelihood_mask_list=likelihood_mask_list)
num_bands = len(self.kwargs_data_joint['multi_band_list'])
if num_bands > 1:
f, axes = plt.subplots(num_bands, 3, figsize=(22, 18))
else:
f, axes = plt.subplots(num_bands, 3, figsize=(22, 6))
for band_index in range(num_bands):
if num_bands > 1:
ax1 = axes[band_index, 0]
ax2 = axes[band_index, 1]
ax3 = axes[band_index, 2]
img_index = band_index
else:
ax1 = axes[0]
ax2 = axes[1]
ax3 = axes[2]
img_index = data_index
model_plot.data_plot(ax=ax1,
band_index=band_index,
text='Observed' + text,
font_size=font_size)
model_plot.model_plot(ax=ax2,
image_names=True,
band_index=band_index,
font_size=font_size,
text='Modeled' + text)
model_plot.normalized_residual_plot(ax=ax3,
v_min=-6,
v_max=6,
band_index=band_index,
font_size=font_size)
f.savefig(img_name + 'residual.pdf', bbox_inches='tight')
if fig_close:
plt.close(f)
if numPix_s is None:
numPix_s = self.kwargs_data_joint['multi_band_list'][0][0][
'image_data'].shape[0]
f_s, axes_s = plt.subplots(1, 1, figsize=(9, 6))
model_plot.source_plot(ax=axes_s,
deltaPix_source=deltaPix_s,
numPix=numPix_s,
center=center,
band_index=band_index,
scale_size=scale_size,
font_size=font_size,
text="Source" + text_source,
plot_scale='log',
v_min=-5,
with_caustics=True)
f_s.savefig(img_name + 'source.pdf')
if fig_close:
plt.close(f_s)
'Ddt_sampling': True
}
index = folder.split('idx')[1].split('_')[0]
# save_file = save_pkl_folder+'idx{0}_ID'.format(index)+ID+'_'+savename
save_file = folder[:-1]
multi_band_list, kwargs_model, kwargs_result_best, chain_list, fix_setting, mcmc_new_list = pickle.load(open(save_file,'rb'))
# fixed_lens, fixed_source, fixed_lens_light, fixed_ps, fixed_cosmo = fix_setting
lens_data = np.ones([99,99])
lens_mask = cr_mask(lens_data, 'normal_mask.reg')
framesize = len(multi_band_list[0][0]['image_data']) #81
ct = int((len(lens_data) - framesize)/2)
lens_mask = (1-lens_mask)[ct:-ct,ct:-ct]
labels_new = [r"$\gamma$", r"$D_{\Delta t}$","H$_0$" ]
modelPlot = ModelPlot(multi_band_list, kwargs_model, kwargs_result_best, arrow_size=0.02, cmap_string="gist_heat",
likelihood_mask_list= [lens_mask])
f, axes = modelPlot.plot_main()
plt.show()
# f, axes = modelPlot.plot_separate()_
# f.show()
# f, axes = modelPlot.plot_subtract_from_data_all()
# f.show()
# multi_band_list = fitting_seq.multi_band_list
# kwargs_psf_updated = multi_band_list[0][1]
# f, axes = chain_plot.psf_iteration_compare(kwargs_psf_updated)
# f.show()
# for i in range(len(chain_list)):
# chain_plot.plot_chain_list(chain_list, i)
# plt.close()
truths=[para_s[0][0]['gamma'],TD_distance, 73.907]
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 #
############
# 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=False,
rescale_pixels_type=None,
log_pixels=False,
add_pixel_noise=cfg.data.add_pixel_noise,
eff_exposure_time={"TDLMC_F160W": test_cfg.data.eff_exposure_time},
train_Y_mean=np.zeros((1, len(cfg.data.Y_cols))),
train_Y_std=np.ones((1, len(cfg.data.Y_cols))),
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:
warnings.warn("Destination folder already exists.")
################
# Compile data #
################
# Image data
with torch.no_grad():
for X_, Y_ in test_loader:
X = X_.to(device)
break
X = X.detach().cpu().numpy()
#############
# MCMC loop #
#############
kwargs_lens_eqn_solver = dict(
min_distance=0.05,
search_window=baobab_cfg.instrument['pixel_scale'] *
baobab_cfg.image['num_pix'],
num_iter_max=200)
fm_posterior = ForwardModelingPosterior(
kwargs_lens_eqn_solver=kwargs_lens_eqn_solver,
astrometric_sigma=test_cfg.image_position_likelihood.sigma,
supersampling_factor=baobab_cfg.numerics.supersampling_factor)
# 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."
)
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):
###########################
# Relevant data and prior #
###########################
data_i = master_truth.iloc[lens_i].copy()
lcdm = LCDM(z_lens=data_i['z_lens'],
z_source=data_i['z_src'],
flat=True)
measured_td_wrt0 = np.array(
literal_eval(
realized_time_delays.iloc[lens_i]['measured_td_wrt0']))
n_img = len(measured_td_wrt0) + 1
#print(baobab_cfg.survey_object_dict)
fm_posterior.set_kwargs_data_joint(
image=X[lens_i, 0, :, :],
measured_td=measured_td_wrt0,
measured_td_sigma=test_cfg.time_delay_likelihood.sigma,
survey_object_dict=baobab_cfg.survey_object_dict,
eff_exposure_time=test_cfg.data.eff_exposure_time,
)
# Update solver according to number of lensed images
if test_cfg.numerics.solver_type == 'NONE':
fm_posterior.kwargs_constraints['solver_type'] = 'NONE'
else:
fm_posterior.kwargs_constraints[
'solver_type'] = 'PROFILE_SHEAR' if n_img == 4 else 'ELLIPSE'
fm_posterior.kwargs_constraints['num_point_source_list'] = [n_img]
#print(fm_posterior.kwargs_params['point_source_model'][0][0])
true_D_dt = lcdm.D_dt(H_0=data_i['H0'], Om0=0.3)
# Pull truth param values and initialize walkers there
if test_cfg.numerics.initialize_walkers_to_truth:
fm_posterior.kwargs_lens_init = metadata_utils.get_kwargs_lens_mass(
data_i)
fm_posterior.kwargs_lens_light_init = metadata_utils.get_kwargs_lens_light(
data_i)
fm_posterior.kwargs_source_init = metadata_utils.get_kwargs_src_light(
data_i)
fm_posterior.kwargs_ps_init = metadata_utils.get_kwargs_ps_lensed(
data_i)
fm_posterior.kwargs_special_init = dict(D_dt=true_D_dt)
###########################
# MCMC posterior sampling #
###########################
lens_i_start_time = time.time()
#with script_utils.HiddenPrints():
chain_list_mcmc, kwargs_result_mcmc = fm_posterior.run_mcmc(
test_cfg.numerics.mcmc)
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,
fm_posterior.kwargs_model,
fm_posterior.kwargs_params['lens_model'][2],
fm_posterior.kwargs_params['point_source_model'][2],
fm_posterior.kwargs_params['source_model'][2],
fm_posterior.kwargs_params['special'][2],
fm_posterior.kwargs_constraints,
kwargs_fixed_lens_light=fm_posterior.
kwargs_params['lens_light_model'][2],
verbose=False)
#from lenstronomy.Plots import chain_plot
model_plot = ModelPlot(fm_posterior.multi_band_list,
fm_posterior.kwargs_model,
kwargs_result_mcmc,
arrow_size=0.02,
cmap_string="gist_heat")
plotting_utils.plot_forward_modeling_comparisons(model_plot, out_dir)
# Plot D_dt histogram
D_dt_samples = new_samples_mcmc[
'D_dt'].values # may contain negative values
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 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 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 plot_final_qso_fit(self,
if_annuli=False,
show_plot=True,
arrows=False,
save_plot=False,
target_ID=None):
"""
Plot the compact fitting result, if a QSO is fitted.
"""
data = self.fitting_specify_class.kwargs_data['image_data']
if 'psf_error_map' in self.fitting_specify_class.kwargs_psf.keys():
modelPlot = ModelPlot(
self.fitting_specify_class.
kwargs_data_joint['multi_band_list'],
self.fitting_specify_class.kwargs_model,
self.kwargs_result,
arrow_size=0.02,
cmap_string="gist_heat",
likelihood_mask_list=self.fitting_specify_class.
kwargs_likelihood['image_likelihood_mask_list'])
_, psf_error_map, _, _ = modelPlot._imageModel.image_linear_solve(
inv_bool=True, **self.kwargs_result)
noise = np.sqrt(
self.fitting_specify_class.kwargs_data['noise_map']**2 +
np.abs(psf_error_map[0]))
else:
noise = self.fitting_specify_class.kwargs_data['noise_map']
ps_list = self.image_ps_list
ps_image = np.zeros_like(ps_list[0])
if target_ID is None:
target_ID = 'target_ID'
for i in range(len(ps_list)):
ps_image = ps_image + ps_list[i]
galaxy_list = self.image_host_list
galaxy_image = np.zeros_like(data)
for i in range(len(galaxy_list)):
galaxy_image = galaxy_image + galaxy_list[i]
model = ps_image + galaxy_image
data_removePSF = data - ps_image
norm_residual = (data - model) / noise
flux_list_2d = [data, model, data_removePSF, norm_residual]
label_list_2d = [
'data', 'model', 'data-Point Source', 'normalized residual'
]
flux_list_1d = [data, model, ps_image, galaxy_image]
label_list_1d = [
'data', 'model', 'Point Source',
'{0} galaxy(s)'.format(len(galaxy_list))
]
fig = total_compare(flux_list_2d,
label_list_2d,
flux_list_1d,
label_list_1d,
deltaPix=self.fitting_specify_class.deltaPix,
zp=self.zp,
if_annuli=if_annuli,
arrows=arrows,
show_plot=show_plot,
mask_image=self.fitting_specify_class.
kwargs_likelihood['image_likelihood_mask_list'][0],
target_ID=target_ID)
if save_plot == True:
savename = self.savename
fig.savefig(savename + "_qso_final_plot.pdf")
if show_plot == True:
plt.show()
else:
plt.close()
def test_raise(self):
with self.assertRaises(ValueError):
kwargs_data = sim_util.data_configure_simple(numPix=10,
deltaPix=1,
background_rms=1)
#kwargs_data['image_data'] = np.zeros((10, 10))
kwargs_model = {'source_light_model_list': ['GAUSSIAN']}
kwargs_params = {
'kwargs_lens': [],
'kwargs_source': [{
'amp': 1,
'sigma': 1,
'center_x': 0,
'center_y': 0
}],
'kwargs_ps': [],
'kwargs_lens_light': []
}
lensPlot = ModelPlot(
multi_band_list=[[kwargs_data, {
'psf_type': 'NONE'
}, {}]],
kwargs_model=kwargs_model,
kwargs_params=kwargs_params,
arrow_size=0.02,
cmap_string="gist_heat")
with self.assertRaises(ValueError):
kwargs_data = sim_util.data_configure_simple(numPix=10,
deltaPix=1,
background_rms=1)
# kwargs_data['image_data'] = np.zeros((10, 10))
kwargs_model = {'source_light_model_list': ['GAUSSIAN']}
kwargs_params = {
'kwargs_lens': [],
'kwargs_source': [{
'amp': 1,
'sigma': 1,
'center_x': 0,
'center_y': 0
}],
'kwargs_ps': [],
'kwargs_lens_light': []
}
lensPlot = ModelPlot(multi_band_list=[[kwargs_data, {}, {}]],
kwargs_model=kwargs_model,
kwargs_params=kwargs_params,
arrow_size=0.02,
cmap_string="gist_heat")
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.source_plot(ax=ax,
numPix=10,
deltaPix_source=0.1,
v_min=None,
v_max=None,
with_caustics=False,
caustic_color='yellow',
fsize=15,
plot_scale='bad')
plt.close()
with self.assertRaises(ValueError):
kwargs_data = sim_util.data_configure_simple(numPix=10,
deltaPix=1,
background_rms=1)
# kwargs_data['image_data'] = np.zeros((10, 10))
kwargs_model = {'source_light_model_list': ['GAUSSIAN']}
kwargs_params = {
'kwargs_lens': [],
'kwargs_source': [{
'amp': 1,
'sigma': 1,
'center_x': 0,
'center_y': 0
}],
'kwargs_ps': [],
'kwargs_lens_light': []
}
lensPlot = ModelPlot(
multi_band_list=[[kwargs_data, {
'psf_type': 'NONE'
}, {}]],
kwargs_model=kwargs_model,
kwargs_params=kwargs_params,
bands_compute=[False],
arrow_size=0.02,
cmap_string="gist_heat")
lensPlot._select_band(band_index=0)
with self.assertRaises(ValueError):
kwargs_data = sim_util.data_configure_simple(numPix=10,
deltaPix=1,
background_rms=1,
exposure_time=1)
# kwargs_data['image_data'] = np.zeros((10, 10))
kwargs_model = {'source_light_model_list': ['GAUSSIAN']}
kwargs_params = {
'kwargs_lens': [],
'kwargs_source': [{
'amp': 1,
'sigma': 1,
'center_x': 0,
'center_y': 0
}],
'kwargs_ps': [],
'kwargs_lens_light': []
}
lensPlot = ModelPlot(
multi_band_list=[[kwargs_data, {
'psf_type': 'NONE'
}, {}]],
kwargs_model=kwargs_model,
kwargs_params=kwargs_params,
bands_compute=[True],
arrow_size=0.02,
cmap_string="gist_heat")
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = lensPlot.source_plot(ax=ax,
numPix=10,
deltaPix_source=0.1,
v_min=None,
v_max=None,
with_caustics=False,
caustic_color='yellow',
fsize=15,
plot_scale='wrong')
plt.close()
if count == 0:
QSO_img = multi_band_list[0][0]['image_data']
plt.imshow(QSO_img, origin='low', norm=LogNorm())
for i in range(len(source_result)):
obj_x, obj_y = len(QSO_img) / 2 - source_result[i][
'center_x'] / pix_scale, len(
QSO_img) / 2 + source_result[i]['center_y'] / pix_scale
plt.text(obj_x,
obj_y,
"obj{0}".format(i),
fontsize=15,
color='k')
plt.show()
modelPlot = ModelPlot(multi_band_list,
kwargs_model,
kwargs_result,
arrow_size=0.02,
cmap_string="gist_heat",
likelihood_mask_list=[QSO_msk])
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)
ds = (x0 - kwargs_ps['ra_image'])**2 + (y0 - kwargs_ps['dec_image'])**2
if ds.min() < 0.01:
x_s.append(x[i])
y_s.append(y[i])
y_grid, x_grid = np.indices(
(framesize, framesize)) #with higher resolution 60*6
# for i in range(len(x_s)):
# lens_mask[np.sqrt((y_grid-y_s[i])**2 + (x_grid-x_s[i])**2) <4] = 0
# plt.imshow(multi_band_list[0][0]['noise_map'],origin='lower')
# plt.show()
# plt.imshow(lens_mask,origin='lower')
# plt.show()
modelPlot = ModelPlot(multi_band_list,
kwargs_model,
kwargs_result,
arrow_size=0.02,
likelihood_mask_list=[lens_mask])
logL = modelPlot._imageModel.likelihood_data_given_model(source_marg=False,
linear_prior=None,
**kwargs_result)
n_data = modelPlot._imageModel.num_data_evaluate
chisq = -logL * 2 / n_data
#!!!
kwargs_result['kwargs_lens'][0]['gamma'] = np.median(chain_list[-1][1][:,
0])
if chisq < 0.0:
print(folder[-4:-1], round(np.median(H0_list), 3))
f, axes = modelPlot.plot_main()
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]
}
# delta_fermat_12 = fermat_pot[0] - fermat_pot[2]
gamma = kwargs_result['kwargs_lens'][0]['gamma']
# phi_ext, gamma_ext = kwargs_result['kwargs_lens'][1]['gamma1'], kwargs_result['kwargs_lens'][1]['gamma2']
mcmc_new_list.append([gamma, D_dt, cal_h0(z_l, z_s, D_dt)])
pickle.dump([
multi_band_list, kwargs_model, kwargs_result, chain_list,
fix_setting, mcmc_new_list
], open(folder + savename, 'wb'))
#%%Print fitting result:
multi_band_list, kwargs_model, kwargs_result, chain_list, fix_setting, mcmc_new_list = pickle.load(
open(folder + savename, 'rb'))
fixed_lens, fixed_source, fixed_lens_light, fixed_ps, fixed_cosmo = fix_setting
labels_new = [r"$\gamma$", r"$D_{\Delta t}$", "H$_0$"]
modelPlot = ModelPlot(multi_band_list,
kwargs_model,
kwargs_result,
arrow_size=0.02,
cmap_string="gist_heat")
f, axes = modelPlot.plot_main()
f.show()
# f, axes = modelPlot.plot_separate()
# f.show()
# f, axes = modelPlot.plot_subtract_from_data_all()
# f.show()
for i in range(len(chain_list)):
chain_plot.plot_chain_list(chain_list, i)
plt.show()
truths = [para_s[0][0]['gamma'], TD_distance, 73.907]
plot = corner.corner(
picklename = 'zoutput/' + 'l35_sersicdisk_1.pkl'
result = pickle.load(open(picklename, 'rb'))
best_fit, chain_list_result, trans_paras, material = result
source_result, image_host, ps_result, image_ps, _ = best_fit
chain_list, _ = chain_list_result
if chain_list[-1][0] == 'EMCEE':
sampler_type, samples_mcmc, param_mcmc, dist_mcmc = chain_list[-1]
mcmc_new_list, labels_new, _ = trans_paras
multi_band_list, kwargs_model, kwargs_result, QSO_msk, kwargs_fixed_source, kwargs_fixed_ps, kwargs_constraints, kwargs_numerics, classes = material
#%% Recover the plot
from lenstronomy.Plots.model_plot import ModelPlot
modelPlot = ModelPlot(multi_band_list,
kwargs_model,
kwargs_result,
arrow_size=0.02,
cmap_string="gist_heat",
likelihood_mask_list=[QSO_msk])
# 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()
# plt.show()