__author__ = "sibirrer"
import scipy.optimize
import numpy as np
import copy
import lenstronomy.Util.param_util as param_util
__all__ = ["Solver2Point"]
[docs]
class Solver2Point(object):
"""Class to solve a constraint lens model with two point source positions.
options are:
'CENTER': solves for 'center_x', 'center_y' parameters of the first lens model
'ELLIPSE': solves for 'e1', 'e2' of the first lens (can also be shear)
'SHAPELETS': solves for shapelet coefficients c01, c10
'THETA_E_PHI: solves for Einstein radius of first lens model and shear angle of second model
"""
[docs]
def __init__(self, lensModel, solver_type="CENTER", decoupling=True):
"""
:param lensModel: instance of LensModel class
:param solver_type: string
:param decoupling: bool
"""
self.lensModel = lensModel
self._lens_mode_list = lensModel.lens_model_list
if solver_type not in [
"CENTER",
"ELLIPSE",
"SHAPELETS",
"THETA_E_PHI",
"THETA_E_ELLIPSE",
]:
raise ValueError("solver_type %s is not a valid option!" % solver_type)
if solver_type == "SHAPELETS":
if not self._lens_mode_list[0] in ["SHAPELETS_CART", "SHAPELETS_POLAR"]:
raise ValueError(
"solver_type %s needs the first lens model to be in ['SHAPELETS_CART', "
"'SHAPELETS_POLAR']" % solver_type
)
if solver_type == "THETA_E_PHI":
if not self._lens_mode_list[1] == "SHEAR":
raise ValueError(
"solver_type %s needs the second lens model to be 'SHEAR"
% solver_type
)
self._solver_type = solver_type
if (
lensModel.multi_plane is True
or "FOREGROUND_SHEAR" in self._lens_mode_list
or solver_type == "THETA_E_PHI"
):
self._decoupling = False
else:
self._decoupling = decoupling
[docs]
def constraint_lensmodel(self, x_pos, y_pos, kwargs_list, xtol=1.49012e-12):
"""Constrains lens model parameters by demanding the solution to match the image
positions to a single source position.
:param x_pos: list of image positions (x-axis)
:param y_pos: list of image position (y-axis)
:param kwargs_list: list of lens model kwargs
:param xtol: tolerance level of solution when to stop the non-linear solver
:return: updated lens model that satisfies the lens equation for the point
sources
"""
kwargs = copy.deepcopy(kwargs_list)
init = self._extract_array(kwargs)
if self._decoupling:
alpha_0_x, alpha_0_y = self.lensModel.alpha(x_pos, y_pos, kwargs)
alpha_1_x, alpha_1_y = self.lensModel.alpha(x_pos, y_pos, kwargs, k=0)
x_sub = alpha_1_x - alpha_0_x
y_sub = alpha_1_y - alpha_0_y
else:
x_sub, y_sub = np.zeros(2), np.zeros(2)
a = self._subtract_constraint(x_sub, y_sub)
x = self.solve(x_pos, y_pos, init, kwargs, a, xtol=xtol)
kwargs = self._update_kwargs(x, kwargs)
y_end = self._F(x, x_pos, y_pos, kwargs, a)
accuracy = np.sum(y_end**2)
return kwargs, accuracy
[docs]
def solve(self, x_pos, y_pos, init, kwargs_list, a, xtol=1.49012e-12):
x = scipy.optimize.fsolve(
self._F, init, args=(x_pos, y_pos, kwargs_list, a), xtol=xtol
) # , factor=0.1)
return x
def _F(self, x, x_pos, y_pos, kwargs_list, a=np.zeros(2)):
kwargs_list = self._update_kwargs(x, kwargs_list)
if self._decoupling:
beta_x, beta_y = self.lensModel.ray_shooting(x_pos, y_pos, kwargs_list, k=0)
else:
beta_x, beta_y = self.lensModel.ray_shooting(x_pos, y_pos, kwargs_list)
y = np.zeros(2)
y[0] = beta_x[0] - beta_x[1]
y[1] = beta_y[0] - beta_y[1]
return y - a
@staticmethod
def _subtract_constraint(x_sub, y_sub):
"""
:param x_sub:
:param y_sub:
:return:
"""
a = np.zeros(2)
a[0] = -x_sub[0] + x_sub[1]
a[1] = -y_sub[0] + y_sub[1]
return a
def _update_kwargs(self, x, kwargs_list):
"""
:param x: list of parameters corresponding to the free parameter of the first lens model in the list
:param kwargs_list: list of lens model kwargs
:return: updated kwargs_list
"""
if self._solver_type == "CENTER":
[center_x, center_y] = x
kwargs_list[0]["center_x"] = center_x
kwargs_list[0]["center_y"] = center_y
elif self._solver_type == "ELLIPSE":
[e1, e2] = x
kwargs_list[0]["e1"] = e1
kwargs_list[0]["e2"] = e2
elif self._solver_type == "SHAPELETS":
[c10, c01] = x
coeffs = list(kwargs_list[0]["coeffs"])
coeffs[1:3] = [c10, c01]
kwargs_list[0]["coeffs"] = coeffs
elif self._solver_type == "THETA_E_PHI":
[theta_E, phi_G] = x
kwargs_list[0]["theta_E"] = theta_E
phi_G_no_sense, gamma_ext = param_util.shear_cartesian2polar(
kwargs_list[1]["gamma1"], kwargs_list[1]["gamma2"]
)
gamma1, gamma2 = param_util.shear_polar2cartesian(phi_G, gamma_ext)
kwargs_list[1]["gamma1"] = gamma1
kwargs_list[1]["gamma2"] = gamma2
elif self._solver_type == "THETA_E_ELLIPSE":
[theta_E, phi_G] = x
kwargs_list[0]["theta_E"] = theta_E
phi_G_no_sense, q = param_util.ellipticity2phi_q(
kwargs_list[0]["e1"], kwargs_list[0]["e2"]
)
e1, e2 = param_util.phi_q2_ellipticity(phi_G, q)
kwargs_list[0]["e1"] = e1
kwargs_list[0]["e2"] = e2
else:
raise ValueError(
"Solver type %s not supported for 2-point solver!" % self._solver_type
)
return kwargs_list
def _extract_array(self, kwargs_list):
"""Inverse of _update_kwargs.
:param kwargs_list:
:return:
"""
if self._solver_type == "CENTER":
center_x = kwargs_list[0]["center_x"]
center_y = kwargs_list[0]["center_y"]
x = [center_x, center_y]
elif self._solver_type == "ELLIPSE":
e1 = kwargs_list[0]["e1"]
e2 = kwargs_list[0]["e2"]
x = [e1, e2]
elif self._solver_type == "SHAPELETS":
coeffs = list(kwargs_list[0]["coeffs"])
[c10, c01] = coeffs[1:3]
x = [c10, c01]
elif self._solver_type == "THETA_E_PHI":
theta_E = kwargs_list[0]["theta_E"]
gamma1 = kwargs_list[1]["gamma1"]
gamma2 = kwargs_list[1]["gamma2"]
phi_ext, gamma_ext = param_util.shear_cartesian2polar(gamma1, gamma2)
x = [theta_E, phi_ext]
elif self._solver_type == "THETA_E_ELLIPSE":
theta_E = kwargs_list[0]["theta_E"]
e1 = kwargs_list[0]["e1"]
e2 = kwargs_list[0]["e2"]
phi_ext, gamma_ext = param_util.shear_cartesian2polar(e1, e2)
x = [theta_E, phi_ext]
else:
raise ValueError(
"Solver type %s not supported for 2-point solver!" % self._solver_type
)
return x
[docs]
def add_fixed_lens(self, kwargs_fixed_lens_list, kwargs_lens_init):
"""
:param kwargs_fixed_lens_list:
:param kwargs_lens_init:
:return:
"""
kwargs_fixed = kwargs_fixed_lens_list[0]
kwargs_lens = kwargs_lens_init[0]
if self._solver_type in ["CENTER"]:
kwargs_fixed["center_x"] = kwargs_lens["center_x"]
kwargs_fixed["center_y"] = kwargs_lens["center_y"]
elif self._solver_type in ["ELLIPSE"]:
kwargs_fixed["e1"] = kwargs_lens["e1"]
kwargs_fixed["e2"] = kwargs_lens["e2"]
elif self._solver_type == "SHAPELETS":
pass
elif self._solver_type == "THETA_E_PHI":
kwargs_fixed["theta_E"] = kwargs_lens["theta_E"]
kwargs_fixed_lens_list[1]["gamma2"] = 0
elif self._solver_type == "THETA_E_ELLIPSE":
kwargs_fixed["theta_E"] = kwargs_lens["theta_E"]
kwargs_fixed_lens_list[0]["e2"] = 0
else:
raise ValueError(
"Solver type %s not supported for 2-point solver!" % self._solver_type
)
return kwargs_fixed_lens_list