def test_plot_line_set(self):
coords = Coordinates(transform_pix2angle=[[1, 0], [0, 1]], ra_at_xy_0=0, dec_at_xy_0=0)
line_set_x = np.linspace(start=0, stop=1, num=10)
line_set_y = np.linspace(start=0, stop=1, num=10)
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = plot_util.plot_line_set(ax, coords, line_set_x, line_set_y, origin=None, color='g', flipped_x=True,
pixel_offset=False)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = plot_util.plot_line_set(ax, coords, line_set_x, line_set_y, origin=[1, 1], color='g', flipped_x=False,
pixel_offset=True)
plt.close()
# and here we input a list of arrays
line_set_list_x = [np.linspace(start=0, stop=1, num=10), np.linspace(start=0, stop=1, num=10)]
line_set_list_y = [np.linspace(start=0, stop=1, num=10), np.linspace(start=0, stop=1, num=10)]
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = plot_util.plot_line_set(ax, coords, line_set_list_x, line_set_list_y, origin=None, color='g',
flipped_x=True)
plt.close()
f, ax = plt.subplots(1, 1, figsize=(4, 4))
ax = plot_util.plot_line_set(ax, coords, line_set_list_x, line_set_list_y, origin=[1, 1], color='g',
flipped_x=False)
plt.close()
def lens_model_plot(ax, lensModel, kwargs_lens, numPix=500, deltaPix=0.01, sourcePos_x=0, sourcePos_y=0,
point_source=False, with_caustics=False, coord_center_ra=0, coord_center_dec=0,
coord_inverse=False):
"""
plots a lens model (convergence) and the critical curves and caustics
:param ax:
:param kwargs_lens:
:param numPix:
:param deltaPix:
:return:
"""
kwargs_data = sim_util.data_configure_simple(numPix, deltaPix, center_ra=coord_center_ra, center_dec=coord_center_dec,
inverse=coord_inverse)
data = ImageData(**kwargs_data)
_coords = data
_frame_size = numPix * deltaPix
x_grid, y_grid = data.pixel_coordinates
lensModelExt = LensModelExtensions(lensModel)
#ra_crit_list, dec_crit_list, ra_caustic_list, dec_caustic_list = lensModelExt.critical_curve_caustics(
# kwargs_lens, compute_window=_frame_size, grid_scale=deltaPix/2.)
x_grid1d = util.image2array(x_grid)
y_grid1d = util.image2array(y_grid)
kappa_result = lensModel.kappa(x_grid1d, y_grid1d, kwargs_lens)
kappa_result = util.array2image(kappa_result)
im = ax.matshow(np.log10(kappa_result), origin='lower', extent=[0, _frame_size, 0, _frame_size], cmap='Greys',
vmin=-1, vmax=1) #, cmap=self._cmap, vmin=v_min, vmax=v_max)
if with_caustics is True:
ra_crit_list, dec_crit_list = lensModelExt.critical_curve_tiling(kwargs_lens, compute_window=_frame_size,
start_scale=deltaPix, max_order=10)
ra_caustic_list, dec_caustic_list = lensModel.ray_shooting(ra_crit_list, dec_crit_list, kwargs_lens)
plot_util.plot_line_set(ax, _coords, ra_caustic_list, dec_caustic_list, color='g')
plot_util.plot_line_set(ax, _coords, ra_crit_list, dec_crit_list, color='r')
if point_source:
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
solver = LensEquationSolver(lensModel)
theta_x, theta_y = solver.image_position_from_source(sourcePos_x, sourcePos_y, kwargs_lens,
min_distance=deltaPix, search_window=deltaPix*numPix)
mag_images = lensModel.magnification(theta_x, theta_y, kwargs_lens)
x_image, y_image = _coords.map_coord2pix(theta_x, theta_y)
abc_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K']
for i in range(len(x_image)):
x_ = (x_image[i] + 0.5) * deltaPix
y_ = (y_image[i] + 0.5) * deltaPix
ax.plot(x_, y_, 'dk', markersize=4*(1 + np.log(np.abs(mag_images[i]))), alpha=0.5)
ax.text(x_, y_, abc_list[i], fontsize=20, color='k')
x_source, y_source = _coords.map_coord2pix(sourcePos_x, sourcePos_y)
ax.plot((x_source + 0.5) * deltaPix, (y_source + 0.5) * deltaPix, '*k', markersize=10)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.autoscale(False)
return ax
def error_map_source_plot(self, ax, numPix, deltaPix_source, v_min=None, v_max=None, with_caustics=False,
font_size=15, point_source_position=True):
"""
plots the uncertainty in the surface brightness in the source from the linear inversion by taking the diagonal
elements of the covariance matrix of the inversion of the basis set to be propagated to the source plane.
#TODO illustration of the uncertainties in real space with the full covariance matrix is subtle. The best way is probably to draw realizations from the covariance matrix.
:param ax: matplotlib axis instance
:param numPix: number of pixels in plot per axis
:param deltaPix_source: pixel spacing in the source resolution illustrated in plot
:param v_min: minimum plotting scale of the map
:param v_max: maximum plotting scale of the map
:param with_caustics: plot the caustics on top of the source reconstruction (may take some time)
:param font_size: font size of labels
:param point_source_position: boolean, if True, plots a point at the position of the point source
:return: plot of source surface brightness errors in the reconstruction on the axis instance
"""
x_grid_source, y_grid_source = util.make_grid_transformed(numPix,
self._coords.transform_pix2angle * deltaPix_source / self._deltaPix)
x_center = self._kwargs_source_partial[0]['center_x']
y_center = self._kwargs_source_partial[0]['center_y']
x_grid_source += x_center
y_grid_source += y_center
coords_source = Coordinates(self._coords.transform_pix2angle * deltaPix_source / self._deltaPix, ra_at_xy_0=x_grid_source[0],
dec_at_xy_0=y_grid_source[0])
error_map_source = self.bandmodel.error_map_source(self._kwargs_source_partial, x_grid_source, y_grid_source,
self._cov_param, model_index_select=False)
error_map_source = util.array2image(error_map_source)
d_s = numPix * deltaPix_source
im = ax.matshow(error_map_source, origin='lower', extent=[0, d_s, 0, d_s],
cmap=self._cmap, vmin=v_min, vmax=v_max) # source
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.autoscale(False)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
cb = plt.colorbar(im, cax=cax)
cb.set_label(r'error variance', fontsize=font_size)
if with_caustics:
ra_caustic_list, dec_caustic_list = self._caustics()
plot_util.plot_line_set(ax, coords_source, ra_caustic_list, dec_caustic_list, color='b')
plot_util.scale_bar(ax, d_s, dist=0.1, text='0.1"', color='w', flipped=False, font_size=font_size)
plot_util.coordinate_arrows(ax, d_s, coords_source,
arrow_size=self._arrow_size, color='w', font_size=font_size)
plot_util.text_description(ax, d_s, text="Error map in source", color="w",
backgroundcolor='k', flipped=False, font_size=font_size)
if point_source_position is True:
ra_source, dec_source = self.bandmodel.PointSource.source_position(self._kwargs_ps_partial, self._kwargs_lens_partial)
plot_util.source_position_plot(ax, coords_source, ra_source, dec_source)
return ax
def deflection_plot(self, ax, v_min=None, v_max=None, axis=0,
with_caustics=False, image_name_list=None,
text="Deflection model", font_size=15,
colorbar_label=r'arcsec'):
"""
:param kwargs_lens:
:param kwargs_else:
:return:
"""
alpha1, alpha2 = self._lensModel.alpha(self._x_grid, self._y_grid, self._kwargs_lens_partial)
alpha1 = util.array2image(alpha1)
alpha2 = util.array2image(alpha2)
if axis == 0:
alpha = alpha1
else:
alpha = alpha2
im = ax.matshow(alpha, origin='lower', extent=[0, self._frame_size, 0, self._frame_size],
vmin=v_min, vmax=v_max, cmap=self._cmap, alpha=0.5)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.autoscale(False)
plot_util.scale_bar(ax, self._frame_size, dist=1, text='1"', color='k', font_size=font_size)
plot_util.coordinate_arrows(ax, self._frame_size, self._coords, color='k',
arrow_size=self._arrow_size, font_size=font_size)
plot_util.text_description(ax, self._frame_size, text=text, color="k",
backgroundcolor='w', font_size=font_size)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
cb = plt.colorbar(im, cax=cax)
cb.set_label(colorbar_label, fontsize=font_size)
if with_caustics is True:
ra_crit_list, dec_crit_list = self._critical_curves()
ra_caustic_list, dec_caustic_list = self._caustics()
plot_util.plot_line_set(ax, self._coords, ra_caustic_list, dec_caustic_list, color='b')
plot_util.plot_line_set(ax, self._coords, ra_crit_list, dec_crit_list, color='r')
ra_image, dec_image = self.bandmodel.PointSource.image_position(self._kwargs_ps_partial, self._kwargs_lens_partial)
plot_util.image_position_plot(ax, self._coords, ra_image, dec_image, image_name_list=image_name_list)
return ax
def arrival_time_surface(ax,
lensModel,
kwargs_lens,
numPix=500,
deltaPix=0.01,
sourcePos_x=0,
sourcePos_y=0,
with_caustics=False,
point_source=False,
n_levels=10,
kwargs_contours={},
image_color_list=None,
letter_font_size=20):
"""
:param ax:
:param lensModel:
:param kwargs_lens:
:param numPix:
:param deltaPix:
:param sourcePos_x:
:param sourcePos_y:
:param with_caustics:
:return:
"""
kwargs_data = sim_util.data_configure_simple(numPix, deltaPix)
data = ImageData(**kwargs_data)
_frame_size = numPix * deltaPix
_coords = data
x_grid, y_grid = data.pixel_coordinates
lensModelExt = LensModelExtensions(lensModel)
#ra_crit_list, dec_crit_list, ra_caustic_list, dec_caustic_list = lensModelExt.critical_curve_caustics(
# kwargs_lens, compute_window=_frame_size, grid_scale=deltaPix/2.)
x_grid1d = util.image2array(x_grid)
y_grid1d = util.image2array(y_grid)
fermat_surface = lensModel.fermat_potential(x_grid1d, y_grid1d,
kwargs_lens, sourcePos_x,
sourcePos_y)
fermat_surface = util.array2image(fermat_surface)
#, cmap='Greys', vmin=-1, vmax=1) #, cmap=self._cmap, vmin=v_min, vmax=v_max)
if with_caustics is True:
ra_crit_list, dec_crit_list = lensModelExt.critical_curve_tiling(
kwargs_lens,
compute_window=_frame_size,
start_scale=deltaPix / 5,
max_order=10)
ra_caustic_list, dec_caustic_list = lensModel.ray_shooting(
ra_crit_list, dec_crit_list, kwargs_lens)
plot_util.plot_line_set(ax,
_coords,
ra_caustic_list,
dec_caustic_list,
shift=_frame_size / 2.,
color='g')
plot_util.plot_line_set(ax,
_coords,
ra_crit_list,
dec_crit_list,
shift=_frame_size / 2.,
color='r')
if point_source is True:
from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
solver = LensEquationSolver(lensModel)
theta_x, theta_y = solver.image_position_from_source(
sourcePos_x,
sourcePos_y,
kwargs_lens,
min_distance=deltaPix,
search_window=deltaPix * numPix)
fermat_pot_images = lensModel.fermat_potential(theta_x, theta_y,
kwargs_lens)
im = ax.contour(
x_grid,
y_grid,
fermat_surface,
origin='lower', # extent=[0, _frame_size, 0, _frame_size],
levels=np.sort(fermat_pot_images),
**kwargs_contours)
mag_images = lensModel.magnification(theta_x, theta_y, kwargs_lens)
x_image, y_image = _coords.map_coord2pix(theta_x, theta_y)
abc_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K']
for i in range(len(x_image)):
x_ = (x_image[i] + 0.5) * deltaPix - _frame_size / 2
y_ = (y_image[i] + 0.5) * deltaPix - _frame_size / 2
if image_color_list is None:
color = 'k'
else:
color = image_color_list[i]
ax.plot(x_, y_, 'x', markersize=10, alpha=1, color=color
) # markersize=8*(1 + np.log(np.abs(mag_images[i])))
ax.text(x_ + deltaPix,
y_ + deltaPix,
abc_list[i],
fontsize=letter_font_size,
color='k')
x_source, y_source = _coords.map_coord2pix(sourcePos_x, sourcePos_y)
ax.plot((x_source + 0.5) * deltaPix - _frame_size / 2,
(y_source + 0.5) * deltaPix - _frame_size / 2,
'*k',
markersize=20)
else:
vmin = np.min(fermat_surface)
vmax = np.max(fermat_surface)
levels = np.linspace(start=vmin, stop=vmax, num=n_levels)
im = ax.contour(
x_grid,
y_grid,
fermat_surface,
origin='lower', # extent=[0, _frame_size, 0, _frame_size],
levels=levels,
**kwargs_contours)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.autoscale(False)
return ax
def source_plot(self,
ax,
numPix,
deltaPix_source,
center=None,
v_min=None,
v_max=None,
with_caustics=False,
caustic_color='yellow',
font_size=15,
plot_scale='log',
scale_size=0.1,
text="Reconstructed source",
colorbar_label=r'log$_{10}$ flux',
point_source_position=True,
**kwargs):
"""
:param ax:
:param numPix:
:param deltaPix_source:
:param center: [center_x, center_y], if specified, uses this as the center
:param v_min:
:param v_max:
:param caustic_color:
:param font_size:
:param plot_scale: string, log or linear, scale of surface brightness plot
:param kwargs:
:return:
"""
if v_min is None:
v_min = self._v_min_default
if v_max is None:
v_max = self._v_max_default
d_s = numPix * deltaPix_source
source, coords_source = self.source(numPix,
deltaPix_source,
center=center)
if plot_scale == 'log':
source[source < 10**v_min] = 10**(
v_min) # to remove weird shadow in plot
source_scale = np.log10(source)
elif plot_scale == 'linear':
source_scale = source
else:
raise ValueError(
'variable plot_scale needs to be "log" or "linear", not %s.' %
plot_scale)
im = ax.matshow(source_scale,
origin='lower',
extent=[0, d_s, 0, d_s],
cmap=self._cmap,
vmin=v_min,
vmax=v_max) # source
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.autoscale(False)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
cb = plt.colorbar(im, cax=cax)
cb.set_label(colorbar_label, fontsize=font_size)
if with_caustics is True:
ra_caustic_list, dec_caustic_list = self._caustics()
plot_util.plot_line_set(ax,
coords_source,
ra_caustic_list,
dec_caustic_list,
color=caustic_color,
points_only=self._caustic_points_only)
plot_util.plot_line_set(ax,
coords_source,
ra_caustic_list,
dec_caustic_list,
color=caustic_color,
points_only=self._caustic_points_only,
**kwargs.get('kwargs_caustic', {}))
plot_util.scale_bar(ax,
d_s,
dist=scale_size,
text='{:.1f}"'.format(scale_size),
color='w',
flipped=False,
font_size=font_size)
if 'no_arrow' not in kwargs or not kwargs['no_arrow']:
plot_util.coordinate_arrows(ax,
self._frame_size,
self._coords,
color='w',
arrow_size=self._arrow_size,
font_size=font_size)
plot_util.text_description(ax,
d_s,
text=text,
color="w",
backgroundcolor='k',
flipped=False,
font_size=font_size)
if point_source_position is True:
ra_source, dec_source = self._bandmodel.PointSource.source_position(
self._kwargs_ps_partial, self._kwargs_lens_partial)
plot_util.source_position_plot(ax, coords_source, ra_source,
dec_source)
return ax
def error_map_source_plot(self,
ax,
numPix,
deltaPix_source,
v_min=None,
v_max=None,
with_caustics=False,
font_size=15,
point_source_position=True):
x_grid_source, y_grid_source = util.make_grid_transformed(
numPix, self._coords.transform_pix2angle * deltaPix_source /
self._deltaPix)
x_center = self._kwargs_source_partial[0]['center_x']
y_center = self._kwargs_source_partial[0]['center_y']
x_grid_source += x_center
y_grid_source += y_center
coords_source = Coordinates(self._coords.transform_pix2angle *
deltaPix_source / self._deltaPix,
ra_at_xy_0=x_grid_source[0],
dec_at_xy_0=y_grid_source[0])
error_map_source = self.bandmodel.error_map_source(
self._kwargs_source_partial, x_grid_source, y_grid_source,
self._cov_param)
error_map_source = util.array2image(error_map_source)
d_s = numPix * deltaPix_source
im = ax.matshow(error_map_source,
origin='lower',
extent=[0, d_s, 0, d_s],
cmap=self._cmap,
vmin=v_min,
vmax=v_max) # source
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.autoscale(False)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
cb = plt.colorbar(im, cax=cax)
cb.set_label(r'error variance', fontsize=font_size)
if with_caustics:
ra_caustic_list, dec_caustic_list = self._caustics()
plot_util.plot_line_set(ax,
coords_source,
ra_caustic_list,
dec_caustic_list,
color='b')
plot_util.scale_bar(ax,
d_s,
dist=0.1,
text='0.1"',
color='w',
flipped=False,
font_size=font_size)
plot_util.coordinate_arrows(ax,
d_s,
coords_source,
arrow_size=self._arrow_size,
color='w',
font_size=font_size)
plot_util.text_description(ax,
d_s,
text="Error map in source",
color="w",
backgroundcolor='k',
flipped=False,
font_size=font_size)
if point_source_position is True:
ra_source, dec_source = self.bandmodel.PointSource.source_position(
self._kwargs_ps_partial, self._kwargs_lens_partial)
plot_util.source_position_plot(ax, coords_source, ra_source,
dec_source)
return ax
def caustics_plot(ax,
pixel_grid,
lens_model,
kwargs_lens,
fast_caustic=True,
coord_inverse=False,
color_crit='r',
color_caustic='g',
pixel_offset=False,
*args,
**kwargs):
"""
:param ax: matplotlib axis instance
:param pixel_grid: lenstronomy PixelGrid() instance (or class with inheritance of PixelGrid()
:param lens_model: LensModel() class instance
:param kwargs_lens: lens model keyword argument list
:param fast_caustic: boolean, if True, uses faster but less precise caustic calculation
(might have troubles for the outer caustic (inner critical curve)
:param coord_inverse: bool, if True, inverts the x-coordinates to go from right-to-left
(effectively the RA definition)
:param color_crit: string, color of critical curve
:param color_caustic: string, color of caustic curve
:param pixel_offset: boolean; if True (default plotting), the coordinates are shifted a half a pixel to match with
the matshow() command to center the coordinates in the pixel center
:param args: argument for plotting curve
:param kwargs: keyword arguments for plotting curves
:return: updated matplotlib axis instance
"""
lens_model_ext = LensModelExtensions(lens_model)
pixel_width = pixel_grid.pixel_width
frame_size = np.max(pixel_grid.width)
coord_center_ra, coord_center_dec = pixel_grid.center
ra0, dec0 = pixel_grid.radec_at_xy_0
origin = [ra0, dec0]
if fast_caustic:
ra_crit_list, dec_crit_list, ra_caustic_list, dec_caustic_list = lens_model_ext.critical_curve_caustics(
kwargs_lens,
compute_window=frame_size,
grid_scale=pixel_width,
center_x=coord_center_ra,
center_y=coord_center_dec)
points_only = False
else:
# only supports individual points due to output of critical_curve_tiling definition
points_only = True
ra_crit_list, dec_crit_list = lens_model_ext.critical_curve_tiling(
kwargs_lens,
compute_window=frame_size,
start_scale=pixel_width,
max_order=10,
center_x=coord_center_ra,
center_y=coord_center_dec)
ra_caustic_list, dec_caustic_list = lens_model.ray_shooting(
ra_crit_list, dec_crit_list, kwargs_lens)
# ra_crit_list, dec_crit_list = list(ra_crit_list), list(dec_crit_list)
# ra_caustic_list, dec_caustic_list = list(ra_caustic_list), list(dec_caustic_list)
plot_util.plot_line_set(ax,
pixel_grid,
ra_caustic_list,
dec_caustic_list,
color=color_caustic,
origin=origin,
flipped_x=coord_inverse,
points_only=points_only,
pixel_offset=pixel_offset,
*args,
**kwargs)
plot_util.plot_line_set(ax,
pixel_grid,
ra_crit_list,
dec_crit_list,
color=color_crit,
origin=origin,
flipped_x=coord_inverse,
points_only=points_only,
pixel_offset=pixel_offset,
*args,
**kwargs)
return ax
def lens_model_plot_custom(image,
ax,
lensModel,
kwargs_lens,
numPix=500,
deltaPix=0.01,
sourcePos_x=0,
sourcePos_y=0,
point_source=False,
with_caustics=False):
"""
Overlay the critical curves and caustics over the provided lensed image
Parameters
----------
image : np.array
ax : matplotlib.pyplot.axes
"""
kwargs_data = sim_util.data_configure_simple(numPix, deltaPix)
data = ImageData(**kwargs_data)
_coords = data
_frame_size = numPix * deltaPix
x_grid, y_grid = data.pixel_coordinates
lensModelExt = LensModelExtensions(lensModel)
_ = ax.matshow(image,
origin='lower',
extent=[0, _frame_size, 0, _frame_size])
if with_caustics is True:
ra_crit_list, dec_crit_list = lensModelExt.critical_curve_tiling(
kwargs_lens,
compute_window=_frame_size,
start_scale=deltaPix,
max_order=20)
ra_caustic_list, dec_caustic_list = lensModel.ray_shooting(
ra_crit_list, dec_crit_list, kwargs_lens)
plot_util.plot_line_set(ax,
_coords,
ra_caustic_list,
dec_caustic_list,
color='y')
plot_util.plot_line_set(ax,
_coords,
ra_crit_list,
dec_crit_list,
color='r')
if point_source:
solver = LensEquationSolver(lensModel)
theta_x, theta_y = solver.image_position_from_source(
sourcePos_x,
sourcePos_y,
kwargs_lens,
min_distance=deltaPix,
search_window=deltaPix * numPix)
mag_images = lensModel.magnification(theta_x, theta_y, kwargs_lens)
x_image, y_image = _coords.map_coord2pix(theta_x, theta_y)
abc_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K']
for i in range(len(x_image)):
x_ = (x_image[i] + 0.5) * deltaPix
y_ = (y_image[i] + 0.5) * deltaPix
ax.plot(x_,
y_,
'dk',
markersize=4 * (1 + np.log(np.abs(mag_images[i]))),
alpha=0.5)
ax.text(x_, y_, abc_list[i], fontsize=20, color='k')
x_source, y_source = _coords.map_coord2pix(sourcePos_x, sourcePos_y)
ax.plot((x_source + 0.5) * deltaPix, (y_source + 0.5) * deltaPix,
'*r',
markersize=5)
#ax.plot(numPix * deltaPix*0.5 + pred['lens_mass_center_x'] + pred['src_light_center_x'], numPix * deltaPix*0.5 + pred['lens_mass_center_y'] + pred['src_light_center_y'], '*k', markersize=5)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
ax.autoscale(False)
return ax
