Fix errors in preprocessing routine

hierarc/LensPosterior/kin_constraints.py

@@ -35,7 +35,7 @@ def __init__(
         num_psf_sampling=100,
         num_kin_sampling=1000,
         multi_observations=False,
-        cosmo_fiducial=None,
+        cosmo_fiducial=None
     ):
         """
 
@@ -110,7 +110,7 @@ def __init__(
             num_psf_sampling=num_psf_sampling,
             num_kin_sampling=num_kin_sampling,
             multi_observations=multi_observations,
-            cosmo_fiducial=cosmo_fiducial,
+            cosmo_fiducial=cosmo_fiducial
         )
 
     def j_kin_draw(self, kwargs_anisotropy, no_error=False):

hierarc/LensPosterior/kin_constraints_composite.py

@@ -3,6 +3,8 @@
 import numpy as np
 from hierarc.LensPosterior.kin_constraints import KinConstraints
 from lenstronomy.Util import constants as const
+from lenstronomy.Analysis.light_profile import LightProfileAnalysis
+from lenstronomy.LightModel.light_model import LightModel
 
 
 class KinConstraintsComposite(KinConstraints):
@@ -12,8 +14,8 @@ def __init__(
         z_source,
         gamma_in_array,
         m2l_array,
-        rho_s_array,
-        r_scale_array,
+        rho0_array,
+        r_s_array,
         theta_E,
         theta_E_error,
         gamma,
@@ -25,16 +27,12 @@ def __init__(
         kwargs_seeing,
         kwargs_numerics_galkin,
         anisotropy_model,
-        kwargs_lens_stars,
         sigma_v_error_independent=None,
         sigma_v_error_covariant=None,
         sigma_v_error_cov_matrix=None,
         kwargs_lens_light=None,
         lens_light_model_list=["HERNQUIST"],
-        lens_model_list=None,
-        MGE_light=False,
         kwargs_mge_light=None,
-        hernquist_approx=True,
         sampling_number=1000,
         num_psf_sampling=100,
         num_kin_sampling=1000,
@@ -46,8 +44,8 @@ def __init__(
         :param z_source: source redshift
         :param gamma_in_array: array of power-law slopes of the mass model
         :param m2l_array: array of mass-to-light ratios of the stellar component
-        :param rho_s_array: array of halo mass normalizations in M_sun / Mpc^3
-        :param r_scale_array: array of halo scale radii in arc seconds
+        :param rho0_array: array of halo mass normalizations in M_sun / Mpc^3
+        :param r_s_array: array of halo scale radii in arc seconds
         :param theta_E: Einstein radius (in arc seconds)
         :param theta_E_error: 1-sigma error on Einstein radius
         :param gamma: power-law slope (2 = isothermal) estimated from imaging data
@@ -73,23 +71,31 @@ def __init__(
             line-of-sight velocity dispersion
         :param anisotropy_model: type of stellar anisotropy model. See details in
             MamonLokasAnisotropy() class of lenstronomy.GalKin.anisotropy
-        :param kwargs_lens_stars: keyword argument list of for stellar mass in the
-            lens model
         :param kwargs_lens_light: keyword argument list of lens light model (optional)
         :param kwargs_mge_light: keyword arguments that go into the MGE decomposition
             routine
-        :param hernquist_approx: bool, if True, uses the Hernquist approximation for the
-            light profile
         :param multi_observations: bool, if True, interprets kwargs_aperture and
             kwargs_seeing as lists of multiple observations
         """
-        self._rho_s_array = rho_s_array
-        self._r_scale_array = r_scale_array
-        self._kappa_s_array, self._r_scale_angle_array = self.get_kappa_s_r_s_angle(
-            rho_s_array, r_scale_array
+        self._light_profile_analysis = LightProfileAnalysis(
+            light_model=LightModel(light_model_list=lens_light_model_list)
+        )
+
+        (
+            amps,
+            sigmas,
+            center_x,
+            center_y,
+        ) = self._light_profile_analysis.multi_gaussian_decomposition(
+            kwargs_lens_light,
+            r_h=r_eff,
+            **kwargs_mge_light
         )
-        self.gamma_in_array = gamma_in_array
-        self.m2l_array = m2l_array
+
+        lens_light_model_list = ["MULTI_GAUSSIAN"]
+        kwargs_lens_light = [{"amp": amps, "sigma": sigmas}]
+
+        lens_model_list = ["GNFW", "MULTI_GAUSSIAN_KAPPA"]
 
         super(KinConstraintsComposite, self).__init__(
             z_lens,
@@ -111,16 +117,22 @@ def __init__(
             kwargs_lens_light=kwargs_lens_light,
             lens_light_model_list=lens_light_model_list,
             lens_model_list=lens_model_list,
-            MGE_light=MGE_light,
-            kwargs_mge_light=kwargs_mge_light,
-            hernquist_approx=hernquist_approx,
+            MGE_light=False, # set False, as MGE is already done as default
+            kwargs_mge_light=None,
+            hernquist_approx=False,
             sampling_number=sampling_number,
             num_psf_sampling=num_psf_sampling,
             num_kin_sampling=num_kin_sampling,
             multi_observations=multi_observations,
         )
 
-        self._kwargs_lens_stars = kwargs_lens_stars
+        self._rho_s_array = rho0_array
+        self._r_scale_array = r_s_array
+        self._kappa_s_array, self._r_scale_angle_array = self.get_kappa_s_r_s_angle(
+            rho0_array, r_s_array
+        )
+        self.gamma_in_array = gamma_in_array
+        self.m2l_array = m2l_array
 
     def get_kappa_s_r_s_angle(self, rho_s, r_scale):
         """Computes the surface mass density of the NFW halo at the scale radius.
@@ -129,8 +141,8 @@ def get_kappa_s_r_s_angle(self, rho_s, r_scale):
         :param r_scale: halo scale radius in arc seconds
         :return: surface mass density divided by the critical density
         """
-        r_s_angle = r_scale / self._lens_cosmo.dd / const.arcsec  # Rs in arcsec
-        kappa_s = rho_s * r_scale / self._lens_cosmo.sigma_crit
+        r_s_angle = r_scale / self.lensCosmo.dd / const.arcsec  # Rs in arcsec
+        kappa_s = rho_s * r_scale / self.lensCosmo.sigma_crit
 
         return kappa_s, r_s_angle
 
@@ -159,6 +171,29 @@ def draw_lens(self, no_error=False):
 
         return kappa_s_draw, r_scale_angle_draw, r_eff_draw, delta_r_eff
 
+    def model_marginalization(self, num_sample_model=20):
+        """
+
+        :param num_sample_model: number of samples drawn from the lens and light model
+            posterior to compute the dimensionless kinematic component J()
+        :return: J() as array for each measurement prediction, covariance matrix in
+            sqrt(J)
+        """
+        num_data = len(self._sigma_v_measured)
+        j_kin_matrix = np.zeros(
+            (num_sample_model, num_data)
+        )  # matrix that contains the sampled J() distribution
+        for i in range(num_sample_model):
+            j_kin = self.j_kin_draw_composite(self.kwargs_anisotropy_base,
+                                              np.mean(self.gamma_in_array),
+                                              np.mean(self.m2l_array),
+                                              no_error=False)
+            j_kin_matrix[i, :] = j_kin
+
+        error_cov_j_sqrt = np.cov(np.sqrt(j_kin_matrix.T))
+        j_model_list = np.mean(j_kin_matrix, axis=0)
+        return j_model_list, error_cov_j_sqrt
+
     def j_kin_draw_composite(self, kwargs_anisotropy, gamma_in, m2l, no_error=False):
         """One simple sampling realization of the dimensionless kinematics of the model.
 
@@ -173,19 +208,18 @@ def j_kin_draw_composite(self, kwargs_anisotropy, gamma_in, m2l, no_error=False)
             no_error=no_error
         )
 
-        kwargs_lens_stars = copy.deepcopy(self._kwargs_lens_stars)
-        for kwargs in kwargs_lens_stars:
-            kwargs["amp"] *= m2l
+        kwargs_lens_stars = copy.deepcopy(self._kwargs_lens_light[0])
 
-            if "sigma" in kwargs:
-                kwargs["sigma"] *= delta_r_eff
-            elif "Rs" in kwargs:
-                kwargs["Rs"] *= delta_r_eff
-            elif "R_sersic" in kwargs:
-                kwargs["R_sersic"] *= delta_r_eff
+        kwargs_lens_stars["amp"] *= m2l
 
-        kwargs_light = copy.deepcopy(self._kwargs_lens_light)
+        if "sigma" in kwargs_lens_stars:
+            kwargs_lens_stars["sigma"] *= delta_r_eff
+        elif "Rs" in kwargs_lens_stars:
+            kwargs_lens_stars["Rs"] *= delta_r_eff
+        elif "R_sersic" in kwargs_lens_stars:
+            kwargs_lens_stars["R_sersic"] *= delta_r_eff
 
+        kwargs_light = copy.deepcopy(self._kwargs_lens_light)
         for kwargs in kwargs_light:
             if "sigma" in kwargs:
                 kwargs["sigma"] *= delta_r_eff
@@ -210,6 +244,8 @@ def j_kin_draw_composite(self, kwargs_anisotropy, gamma_in, m2l, no_error=False)
             kwargs_lens_light=kwargs_light,
             kwargs_anisotropy=kwargs_anisotropy,
             r_eff=r_eff_draw,
+            theta_E=self._theta_E, # send this to avoid unnecessary recomputation
+            gamma=self._gamma, # send this to avoid unnecessary recomputation
         )
         return j_kin
 
@@ -246,6 +282,17 @@ def hierarchy_configuration(self, num_sample_model=20):
         }
         return kwargs_likelihood
 
+    def anisotropy_scaling(self):
+        """
+
+        :return: anisotropy scaling grid along the axes defined by ani_param_array
+        """
+        j_ani_0 = self.j_kin_draw_composite(self.kwargs_anisotropy_base,
+                                            np.mean(self.gamma_in_array),
+                                            np.mean(self.m2l_array),
+                                            no_error=True)
+        return self._anisotropy_scaling_relative(j_ani_0)
+
     def _anisotropy_scaling_relative(self, j_ani_0):
         """Anisotropy scaling relative to a default J prediction.
 
@@ -289,8 +336,7 @@ def _anisotropy_scaling_relative(self, j_ani_0):
                     (
                         len(self.gamma_in_array),
                         len(self.m2l_array),
-                        len(self.ani_param_array[0]),
-                        len(self.ani_param_array[1]),
+                        len(self.ani_param_array),
                     )
                 )
                 for _ in range(num_data)