Merge pull request lenstronomy#654 from sibirrer/ps_redshift

Make point source and lens equation flexible to enable different redshifts

lenstronomy/Cosmo/lens_cosmo.py

@@ -439,3 +439,23 @@ def sersic_k_eff2m_star(self, k_eff, R_sersic, n_sersic):
         )
         m_star = k_eff * self.sigma_crit_angle * norm_integral
         return m_star
+
+    def beta_double_source_plane(self, z_lens, z_source_1, z_source_2):
+        """Model prediction of ratio of scaled deflection angles.
+
+        :param z_lens: lens redshift
+        :param z_source_1: source_1 redshift
+        :param z_source_2: source_2 redshift
+        :param cosmo: ~astropy.cosmology instance
+        :return: beta
+        """
+        ds1 = self.background.cosmo.angular_diameter_distance(z=z_source_1).value
+        dds1 = self.background.cosmo.angular_diameter_distance_z1z2(
+            z1=z_lens, z2=z_source_1
+        ).value
+        ds2 = self.background.cosmo.angular_diameter_distance(z=z_source_2).value
+        dds2 = self.background.cosmo.angular_diameter_distance_z1z2(
+            z1=z_lens, z2=z_source_2
+        ).value
+        beta = dds1 / ds1 * ds2 / dds2
+        return beta

lenstronomy/LensModel/MultiPlane/multi_plane.py

@@ -58,6 +58,23 @@ def __init__(
         :param distance_ratio_sampling: bool, if True, will use sampled
          distance ratios to update T_ij value in multi-lens plane computation.
         """
+        self.kwargs_class = {
+            "z_source": z_source,
+            "lens_model_list": lens_model_list,
+            "lens_redshift_list": lens_redshift_list,
+            "cosmo": cosmo,
+            "numerical_alpha_class": numerical_alpha_class,
+            "observed_convention_index": observed_convention_index,
+            "ignore_observed_positions": ignore_observed_positions,
+            "z_source_convention": z_source_convention,
+            "z_lens_convention": z_lens_convention,
+            "cosmo_interp": cosmo_interp,
+            "z_interp_stop": z_interp_stop,
+            "num_z_interp": num_z_interp,
+            "kwargs_interp": kwargs_interp,
+            "kwargs_synthesis": kwargs_synthesis,
+            "distance_ratio_sampling": distance_ratio_sampling,
+        }
         if z_source_convention is None:
             z_source_convention = z_source
         if z_interp_stop is None:

lenstronomy/LensModel/Solver/epl_shear_solver.py

@@ -281,8 +281,12 @@ def solve_lenseq_pemd(pos_, kwargs_lens, Nmeas=400, Nmeas_extra=80, **kwargs):
     if pos.ndim > 1 and pos.shape[-1] != 1:
         pos = pos[..., None]
     t = kwargs_lens[0]["gamma"] - 1 if "gamma" in kwargs_lens[0] else 1
-
-    theta_ell, q = ellipticity2phi_q(kwargs_lens[0]["e1"], kwargs_lens[0]["e2"])
+    if "e1" in kwargs_lens[0] and "e2" in kwargs_lens[0]:
+        e1 = kwargs_lens[0]["e1"]
+        e2 = kwargs_lens[0]["e2"]
+    else:
+        e1, e2 = 0, 0
+    theta_ell, q = ellipticity2phi_q(e1, e2)
     b = kwargs_lens[0]["theta_E"] * np.sqrt(q)
     if len(kwargs_lens) > 1:
         gamma = kwargs_lens[1]["gamma1"] + 1j * kwargs_lens[1]["gamma2"]

lenstronomy/LensModel/Solver/lens_equation_solver.py

@@ -9,6 +9,7 @@
 __all__ = ["LensEquationSolver"]
 
 SUPPORTED_LENS_MODELS_ANALYTICAL = (
+    ["SIS", "SHEAR"],
     ["SIE", "SHEAR"],
     ["SIE"],
     ["EPL_NUMBA", "SHEAR"],
@@ -30,6 +31,15 @@ def __init__(self, lensModel):
         """
         self.lensModel = lensModel
 
+    def change_source_redshift(self, z_source=None):
+        """Change source redshift in solver.
+
+        :param z_source:
+        :type z_source: float or None
+        :return: updated lens model instance
+        """
+        self.lensModel.change_source_redshift(z_source=z_source)
+
     def image_position_stochastic(
         self,
         source_x,
@@ -208,8 +218,24 @@ def image_position_analytical(
 
         if lens_model_list not in SUPPORTED_LENS_MODELS_ANALYTICAL:
             raise ValueError(
-                "Only SIE, EPL, EPL_NUMBA (+shear +convergence) supported in the analytical solver for now."
+                "Only SIS (only with shear), SIE, EPL, EPL_NUMBA (+shear +convergence) "
+                "supported in the analytical solver for now."
+            )
+
+        if self.lensModel.type != "SinglePlane":
+            raise ValueError(
+                "lens model type %s not supported for analytical lens equation solver, "
+                "Needs to be SinglePlane." % self.lensModel.type
             )
+
+        # re-scale solutions if source redshift has changed (i.e. alpha_scaling != 1)
+        alpha_scaling = self.lensModel.lens_model.alpha_scaling
+        gamma = kwargs_lens[0]["gamma"] if "gamma" in kwargs_lens[0] else 2
+        kwargs_lens_[0]["theta_E"] *= alpha_scaling ** (1.0 / (gamma - 1))
+        if "SHEAR" in lens_model_list:
+            kwargs_lens_[1]["gamma1"] *= alpha_scaling
+            kwargs_lens_[1]["gamma2"] *= alpha_scaling
+
         x_mins, y_mins = solve_lenseq_pemd((x_, y_), kwargs_lens_, **kwargs_solver)
 
         if arrival_time_sort:

lenstronomy/LensModel/lens_model.py

@@ -67,6 +67,8 @@ def __init__(
         self.lens_model_list = lens_model_list
         self.z_lens = z_lens
         self.z_source = z_source
+        if z_source_convention is None and z_source is not None:
+            z_source_convention = z_source
         self._z_source_convention = z_source_convention
         self.redshift_list = lens_redshift_list
 
@@ -92,9 +94,12 @@ def __init__(
             raise ValueError(
                 "You can only have one model for line-of-sight corrections."
             )
-
+        if z_lens is not None and z_source is not None:
+            self._lensCosmo = LensCosmo(z_lens, z_source, cosmo=cosmo)
         # Multi-plane or single-plane lensing?
         self.multi_plane = multi_plane
+        self._decouple_multi_plane = decouple_multi_plane
+        self._los_effects = los_effects
         if multi_plane is True:
             if lens_redshift_list is None:
                 raise ValueError(
@@ -125,6 +130,7 @@ def __init__(
                     kwargs_synthesis=kwargs_synthesis,
                     **kwargs_multiplane_model
                 )
+                self.type = "MultiPlaneDecoupled"
             else:
                 self.lens_model = MultiPlane(
                     z_source,
@@ -141,6 +147,7 @@ def __init__(
                     kwargs_synthesis=kwargs_synthesis,
                     distance_ratio_sampling=distance_ratio_sampling,
                 )
+                self.type = "MultiPlane"
 
         else:
             if los_effects is True:
@@ -153,6 +160,7 @@ def __init__(
                     kwargs_interp=kwargs_interp,
                     kwargs_synthesis=kwargs_synthesis,
                 )
+                self.type = "SinglePlaneLOS"
             else:
                 self.lens_model = SinglePlane(
                     lens_model_list,
@@ -162,9 +170,20 @@ def __init__(
                     kwargs_interp=kwargs_interp,
                     kwargs_synthesis=kwargs_synthesis,
                 )
-
-        if z_lens is not None and z_source is not None:
-            self._lensCosmo = LensCosmo(z_lens, z_source, cosmo=cosmo)
+                self.type = "SinglePlane"
+                if z_source is not None and z_source_convention is not None:
+                    if z_source != z_source_convention:
+                        if z_lens is None:
+                            raise ValueError(
+                                "a lens redshift needs to provided when z_source != z_source_convention"
+                            )
+                        else:
+                            alpha_scaling = self._lensCosmo.beta_double_source_plane(
+                                z_lens=self.z_lens,
+                                z_source_1=z_source,
+                                z_source_2=self._z_source_convention,
+                            )
+                            self.lens_model.change_redshift_scaling(alpha_scaling)
 
     def info(self):
         """Shows what models are being initialized and what parameters are being
@@ -479,6 +498,45 @@ def set_dynamic(self):
         """
         self.lens_model.set_dynamic()
 
+    def change_source_redshift(self, z_source):
+        """Changes the ray-tracing (and all relevant default calculations) to a
+        different source redshift while preserving the deflection angles to
+        z_source_convention.
+
+        :param z_source: source redshift
+        :return: None
+        """
+        if z_source is None:
+            return 0
+        if z_source == self.z_source:
+            return 0
+        if self.multi_plane is True:
+            if self._decouple_multi_plane:
+                raise NotImplementedError(
+                    "MultiPlaneDecoupled lens model does not support change in source redshift"
+                )
+            else:
+                # TODO: is it possible to not re-initialize it for performance improvements?
+                kwargs_lens_class = self.lens_model.kwargs_class
+                kwargs_lens_class["z_source"] = z_source
+                self.lens_model = MultiPlane(**kwargs_lens_class)
+        else:
+            if self._los_effects is True:
+                raise NotImplementedError(
+                    "SinglePlaneLOS lens model does not support change in source redshift"
+                )
+            else:
+                alpha_scaling = self._lensCosmo.beta_double_source_plane(
+                    z_lens=self.z_lens,
+                    z_source_1=z_source,
+                    z_source_2=self._z_source_convention,
+                )
+                self.lens_model.change_redshift_scaling(alpha_scaling)
+
+        if self.z_lens is not None:
+            self._lensCosmo = LensCosmo(self.z_lens, z_source, cosmo=self.cosmo)
+        self.z_source = z_source
+
     def _deflection_differential(self, x, y, kwargs, k=None, diff=0.00001):
         """
 

lenstronomy/LensModel/single_plane.py

@@ -9,6 +9,37 @@
 class SinglePlane(ProfileListBase):
     """Class to handle an arbitrary list of lens models in a single lensing plane."""
 
+    def __init__(
+        self,
+        lens_model_list,
+        numerical_alpha_class=None,
+        lens_redshift_list=None,
+        z_source_convention=None,
+        kwargs_interp=None,
+        kwargs_synthesis=None,
+        alpha_scaling=1,
+    ):
+        """
+
+        :param lens_model_list: list of strings with lens model names
+        :param numerical_alpha_class: an instance of a custom class for use in NumericalAlpha() lens model
+         deflection angles as a lens model. See the documentation in Profiles.numerical_deflections
+        :param kwargs_interp: interpolation keyword arguments specifying the numerics.
+         See description in the Interpolate() class. Only applicable for 'INTERPOL' and 'INTERPOL_SCALED' models.
+        :param kwargs_synthesis: keyword arguments for the 'SYNTHESIS' lens model, if applicable
+        :param alpha_scaling: scaling factor of deflection angle relative to z_source_convention
+        """
+        self._alpha_scaling = alpha_scaling
+        ProfileListBase.__init__(
+            self,
+            lens_model_list=lens_model_list,
+            numerical_alpha_class=numerical_alpha_class,
+            lens_redshift_list=lens_redshift_list,
+            z_source_convention=z_source_convention,
+            kwargs_interp=kwargs_interp,
+            kwargs_synthesis=kwargs_synthesis,
+        )
+
     def ray_shooting(self, x, y, kwargs, k=None):
         """Maps image to source position (inverse deflection).
 
@@ -68,7 +99,7 @@ def potential(self, x, y, kwargs, k=None):
         for i, func in enumerate(self.func_list):
             if bool_list[i] is True:
                 potential += func.function(x, y, **kwargs[i])
-        return potential
+        return potential * self._alpha_scaling
 
     def alpha(self, x, y, kwargs, k=None):
         """Deflection angles.
@@ -95,7 +126,7 @@ def alpha(self, x, y, kwargs, k=None):
                 f_x += f_x_i
                 f_y += f_y_i
 
-        return f_x, f_y
+        return f_x * self._alpha_scaling, f_y * self._alpha_scaling
 
     def hessian(self, x, y, kwargs, k=None):
         """Hessian matrix.
@@ -129,7 +160,28 @@ def hessian(self, x, y, kwargs, k=None):
                 f_xy += f_xy_i
                 f_yx += f_yx_i
                 f_yy += f_yy_i
-        return f_xx, f_xy, f_yx, f_yy
+        return (
+            f_xx * self._alpha_scaling,
+            f_xy * self._alpha_scaling,
+            f_yx * self._alpha_scaling,
+            f_yy * self._alpha_scaling,
+        )
+
+    def change_redshift_scaling(self, alpha_scaling):
+        """
+
+        :param alpha_scaling: scaling parameter of the reduced deflection angle relative to z_source_convention
+        :return: None
+        """
+        self._alpha_scaling = alpha_scaling
+
+    @property
+    def alpha_scaling(self):
+        """Deflector scaling factor.
+
+        :return: alpha_scaling
+        """
+        return self._alpha_scaling
 
     def mass_3d(self, r, kwargs, bool_list=None):
         """Computes the mass within a 3d sphere of radius r.

lenstronomy/PointSource/Types/base_ps.py

@@ -31,6 +31,7 @@ def __init__(
         :param redshift: redshift of the source, only required for multiple source redshifts
         :type redshift: None or float
         """
+        self._redshift = redshift
         self._lens_model = lens_model
         if index_lens_model_list is not None:
             k_list = []

lenstronomy/PointSource/Types/lensed_position.py

@@ -46,6 +46,7 @@ def image_position(
                 kwargs_lens_eqn_solver = {}
             ra_source, dec_source = self.source_position(kwargs_ps, kwargs_lens)
             # TODO: this solver does not distinguish between different frames/bands with partial lens models
+            self._solver.change_source_redshift(self._redshift)
             ra_image, dec_image = self._solver.image_position_from_source(
                 ra_source,
                 dec_source,
@@ -68,6 +69,7 @@ def source_position(self, kwargs_ps, kwargs_lens=None):
         """
         ra_image = kwargs_ps["ra_image"]
         dec_image = kwargs_ps["dec_image"]
+        self._lens_model.change_source_redshift(self._redshift)
 
         if self.k_list is None:
             x_source, y_source = self._lens_model.ray_shooting(
@@ -109,6 +111,7 @@ def image_amplitude(
             details
         :return: array of image amplitudes
         """
+        self._lens_model.change_source_redshift(self._redshift)
         if self._fixed_magnification:
             if x_pos is not None and y_pos is not None:
                 ra_image, dec_image = x_pos, y_pos
@@ -150,6 +153,7 @@ def source_amplitude(self, kwargs_ps, kwargs_lens=None):
         if self._fixed_magnification:
             source_amp = kwargs_ps["source_amp"]
         else:
+            self._lens_model.change_source_redshift(self._redshift)
             ra_image, dec_image = kwargs_ps["ra_image"], kwargs_ps["dec_image"]
             if self.k_list is None:
                 mag = self._lens_model.magnification(ra_image, dec_image, kwargs_lens)

lenstronomy/PointSource/Types/source_position.py

@@ -43,6 +43,7 @@ def image_position(
         if kwargs_lens_eqn_solver is None:
             kwargs_lens_eqn_solver = {}
         ra_source, dec_source = self.source_position(kwargs_ps)
+        self._solver.change_source_redshift(self._redshift)
         ra_image, dec_image = self._solver.image_position_from_source(
             ra_source,
             dec_source,
@@ -87,6 +88,7 @@ def image_amplitude(
         :return: array of image amplitudes
         """
         if self._fixed_magnification:
+            self._lens_model.change_source_redshift(self._redshift)
             if x_pos is not None and y_pos is not None:
                 ra_image, dec_image = x_pos, y_pos
             else:
@@ -122,6 +124,7 @@ def source_amplitude(self, kwargs_ps, kwargs_lens=None):
         if self._fixed_magnification:
             source_amp = kwargs_ps["source_amp"]
         else:
+            self._lens_model.change_source_redshift(self._redshift)
             ra_image, dec_image = self.image_position(kwargs_ps, kwargs_lens)
             mag = self._lens_model.magnification(ra_image, dec_image, kwargs_lens)
             point_amp = kwargs_ps["point_amp"]