Allow for direct Dd, Ddt sampling

hierarc/Diagnostics/goodness_of_fit.py

@@ -335,7 +335,8 @@ def plot_ifu_fit(
         if len(np.atleast_1d(bin_edges)) < 2:
             bin_edges = np.arange(len(sigma_v_measurement) + 1) * bin_edges
         r_bins = bin_edges[:-1] + np.diff(bin_edges) / 2
-
+        print(r_bins, "r_bins")
+        print(sigma_v_measurement, "sigma_v_measurement")
         ax.errorbar(
             x=r_bins,
             y=sigma_v_measurement,

hierarc/Likelihood/LensLikelihood/base_lens_likelihood.py

@@ -34,6 +34,7 @@ def __init__(
         name="name",
         normalized=False,
         kwargs_lens_properties=None,
+        distance_sampling=None,
         **kwargs_likelihood
     ):
         """
@@ -47,6 +48,7 @@ def __init__(
          (in case of a Gaussian 1/(sigma sqrt(2 pi)) ) to compute the reduced chi2 statistics
         :param kwargs_lens_properties: keyword arguments of the lens properties
         :param kwargs_likelihood: keyword arguments specifying the likelihood function,
+        :param distance_sampling: dictionary specifying which distances are sampled for this lens (Used only if cosmology='FREE')
         see individual classes for their use
         """
         self.name = name

hierarc/Likelihood/cosmo_likelihood.py

@@ -57,6 +57,7 @@ def __init__(
         :param normalized: bool, if True, returns the normalized likelihood, if False, separates the constant prefactor
          (in case of a Gaussian 1/(sigma sqrt(2 pi)) ) to compute the reduced chi2 statistics
         """
+
         self._cosmology = cosmology
         self._kwargs_lens_list = kwargs_likelihood_list
         if kwargs_model.get("sigma_v_systematics", False) is True:
@@ -67,8 +68,15 @@ def __init__(
             kwargs_global_model=kwargs_model,
         )
         gamma_pl_num = self._likelihoodLensSample.gamma_pl_num
+        Ddt_sampling_num = self._likelihoodLensSample.Ddt_sampling_num
+        Dd_sampling_num = self._likelihoodLensSample.Dd_sampling_num
         self.param = ParamManager(
-            cosmology, gamma_pl_num=gamma_pl_num, **kwargs_model, **kwargs_bounds
+            cosmology,
+            gamma_pl_num=gamma_pl_num,
+            Dd_sampling_num=Dd_sampling_num,
+            Ddt_sampling_num=Ddt_sampling_num,
+            **kwargs_model,
+            **kwargs_bounds
         )
         self._lower_limit, self._upper_limit = self.param.param_bounds
         self._prior_add = False
@@ -219,7 +227,10 @@ def cosmo_instance(self, kwargs_cosmo):
                 K=kwargs_cosmo.get("K", None),
             )
             return cosmo
-        if self._cosmo_fixed is None:
+
+        if self._cosmology == "FREE":
+            cosmo = None
+        elif self._cosmo_fixed is None:
             cosmo = self.param.cosmo(kwargs_cosmo)
             if self._interpolate_cosmo is True:
                 cosmo = CosmoInterp(

hierarc/Likelihood/hierarchy_likelihood.py

@@ -48,6 +48,8 @@ def __init__(
         gamma_pl_index=None,
         gamma_pl_global_sampling=False,
         gamma_pl_global_dist="NONE",
+        Ddt_sampling_index=None,
+        Dd_sampling_index=None,
         # kinematic model quantities
         kin_scaling_param_list=None,
         j_kin_scaling_param_axes=None,
@@ -112,6 +114,8 @@ def __init__(
         :param gamma_pl_global_sampling: if sampling a global power-law density slope distribution
         :type gamma_pl_global_sampling: bool
         :param gamma_pl_global_dist: distribution of global gamma_pl distribution ("GAUSSIAN" or "NONE")
+        :param Ddt_sampling_index: index of Ddt parameter associated with this lens (this used if cosmology='FREE', else None)
+        :param Dd_sampling_index: index of Dd parameter associated with this lens (this used if cosmology='FREE', else None)
         :param normalized: bool, if True, returns the normalized likelihood, if False, separates the constant prefactor
          (in case of a Gaussian 1/(sigma sqrt(2 pi)) ) to compute the reduced chi2 statistics
         :param kwargs_lens_properties: keyword arguments of the lens properties
@@ -204,6 +208,9 @@ def __init__(
         else:
             self._inclination_sampling = False
 
+        self._Ddt_sampling_index = Ddt_sampling_index
+        self._Dd_sampling_index = Dd_sampling_index
+
     def info(self):
         """Information about the lens.
 
@@ -242,7 +249,13 @@ def lens_log_likelihood(
         # here we compute the unperturbed angular diameter distances of the lens system given the cosmology
         # Note: Distances are in physical units of Mpc. Make sure the posteriors to evaluate this likelihood is in the
         # same units
-        ddt, dd = self.angular_diameter_distances(cosmo)
+        if cosmo is None:
+            ddt, dd = (
+                kwargs_lens["Ddt_list"][self._Ddt_sampling_index],
+                kwargs_lens["Dd_list"][self._Dd_sampling_index],
+            )
+        else:
+            ddt, dd = self.angular_diameter_distances(cosmo)
         beta_dsp = self.beta_dsp(cosmo)
         kwargs_source = self._kwargs_init(kwargs_source)
         z_apparent_m_anchor = kwargs_source.get("z_apparent_m_anchor", 0.1)

hierarc/Likelihood/lens_sample_likelihood.py

@@ -21,26 +21,46 @@ def __init__(self, kwargs_lens_list, normalized=False, kwargs_global_model=None)
             kwargs_global_model = {}
         self._lens_list = []
         self._gamma_pl_num = 0
+        self._Ddt_num = 0
+        self._Dd_num = 0
         gamma_pl_index = 0
+        Ddt_index = 0
+        Dd_index = 0
 
         gamma_pl_global_sampling = kwargs_global_model.get(
             "gamma_pl_global_sampling", False
         )
 
         for kwargs_lens in kwargs_lens_list:
             gamma_pl_index_ = None
+            Dd_sampling_index_ = None
+            Ddt_sampling_index_ = None
             if "kin_scaling_param_list" in kwargs_lens and not gamma_pl_global_sampling:
                 kin_scaling_param_list = kwargs_lens["kin_scaling_param_list"]
                 if "gamma_pl" in kin_scaling_param_list:
                     self._gamma_pl_num += 1
                     gamma_pl_index_ = copy.deepcopy(gamma_pl_index)
                     gamma_pl_index += 1
+
+            if kwargs_lens.get("distance_sampling", False):
+                if kwargs_lens["distance_sampling"].get("Ddt_sampling", False):
+                    self._Ddt_num += 1
+                    Ddt_sampling_index_ = copy.deepcopy(Ddt_index)
+                    Ddt_index += 1
+                if kwargs_lens["distance_sampling"].get("Dd_sampling", False):
+                    self._Dd_num += 1
+                    Dd_sampling_index_ = copy.deepcopy(Dd_index)
+                    Dd_index += 1
+
             kwargs_lens_ = self._merge_global2local_settings(
                 kwargs_global_model=kwargs_global_model, kwargs_lens=kwargs_lens
             )
+
             self._lens_list.append(
                 LensLikelihood(
                     gamma_pl_index=gamma_pl_index_,
+                    Ddt_sampling_index=Ddt_sampling_index_,
+                    Dd_sampling_index=Ddt_sampling_index_,
                     normalized=normalized,
                     **kwargs_lens_
                 )
@@ -108,6 +128,25 @@ def gamma_pl_num(self):
         """
         return self._gamma_pl_num
 
+    @property
+    def Ddt_sampling_num(self):
+        """Number of time-delay distance parameters being sampled on individual lenses.
+
+        :return: number of time-delay distance parameters being sampled on individual
+            lenses
+        """
+        return self._Ddt_num
+
+    @property
+    def Dd_sampling_num(self):
+        """Number of angular diameter distance to the deflector parameters being sampled
+        on individual lenses.
+
+        :return: number of angular diameter distance to the deflector parameters being
+            sampled on individual lenses
+        """
+        return self._Dd_num
+
     @staticmethod
     def _merge_global2local_settings(kwargs_global_model, kwargs_lens):
         """

hierarc/Sampling/Distributions/lens_distribution.py

@@ -115,6 +115,8 @@ def draw_lens(
         gamma_pl_list=None,
         gamma_pl_mean=2,
         gamma_pl_sigma=0,
+        Ddt_list=None,
+        Dd_list=None,
     ):
         """Draws a realization of a specific model from the hyperparameter distribution.
 

hierarc/Sampling/ParamManager/cosmo_param.py

@@ -29,6 +29,7 @@ def __init__(
             "oLCDM",
             "wphiCDM",
             "NONE",
+            "FREE",
         ]
         if cosmology not in self._supported_cosmologies:
             raise ValueError(
@@ -43,7 +44,7 @@ def param_list(self, latex_style=False):
         :return: list of the free parameters being sampled in the same order as the sampling
         """
         list = []
-        if self._cosmology not in ["NONE"]:
+        if self._cosmology not in ["NONE", "FREE"]:
             if "h0" not in self._kwargs_fixed:
                 if latex_style is True:
                     list.append(r"$H_0$")
@@ -112,7 +113,7 @@ def args2kwargs(self, args, i=0):
         :return: keyword argument list with parameter names
         """
         kwargs = {}
-        if self._cosmology not in ["NONE"]:
+        if self._cosmology not in ["NONE", "FREE"]:
             if "h0" in self._kwargs_fixed:
                 kwargs["h0"] = self._kwargs_fixed["h0"]
             else:
@@ -180,7 +181,7 @@ def kwargs2args(self, kwargs):
         :return: sampling argument list in specified order
         """
         args = []
-        if self._cosmology not in ["NONE"]:
+        if self._cosmology not in ["NONE", "FREE"]:
             if "h0" not in self._kwargs_fixed:
                 args.append(kwargs["h0"])
             if self._cosmology in ["FLCDM", "FwCDM", "w0waCDM", "oLCDM", "wphiCDM"]:
@@ -245,6 +246,8 @@ def cosmo(self, kwargs):
             )
         elif self._cosmology == "NONE":
             cosmo = None
+        elif self._cosmology == "FREE":
+            cosmo = None
         else:
             raise ValueError("Cosmology %s is not supported" % self._cosmology)
         return cosmo

hierarc/Sampling/ParamManager/lens_param.py

@@ -23,6 +23,8 @@ def __init__(
         gamma_pl_global_dist="NONE",
         kwargs_fixed=None,
         log_scatter=False,
+        Ddt_sampling_num=0,
+        Dd_sampling_num=0,
     ):
         """
 
@@ -55,6 +57,8 @@ def __init__(
         :param log_scatter: boolean, if True, samples the Gaussian scatter amplitude in log space (and thus flat prior
          in log)
         :param kwargs_fixed: keyword arguments that are held fixed through the sampling
+        :param Ddt_sampling_num: int, number of time-delay distances to be sampled (to be assigned to individual lenses)
+        :param Dd_sampling_num: int, number of angular diameter distances to be sampled (to be assigned to individual lenses)
         """
         self._lambda_mst_sampling = lambda_mst_sampling
         self._lambda_mst_distribution = lambda_mst_distribution
@@ -71,6 +75,8 @@ def __init__(
         self._gamma_pl_num = gamma_pl_num
         self._gamma_pl_global_sampling = gamma_pl_global_sampling
         self._gamma_pl_global_dist = gamma_pl_global_dist
+        self._Ddt_sampling_num = Ddt_sampling_num
+        self._Dd_sampling_num = Dd_sampling_num
 
         self._log_scatter = log_scatter
         if kwargs_fixed is None:
@@ -185,6 +191,18 @@ def param_list(self, latex_style=False):
                         list.append(r"$\sigma(\gamma_{\rm pl, global})$")
                     else:
                         list.append("gamma_pl_sigma")
+        for i in range(self._Ddt_sampling_num):
+            if "Ddt_%i" % i not in self._kwargs_fixed:
+                if latex_style is True:
+                    list.append(r"$D_{\Delta t, %i}$" % i)
+                else:
+                    list.append("Ddt_%i" % i)
+        for i in range(self._Dd_sampling_num):
+            if "Dd_%i" % i not in self._kwargs_fixed:
+                if latex_style is True:
+                    list.append(r"$D_{\rm d, %i}$" % i)
+                else:
+                    list.append("Dd_%i" % i)
         return list
 
     def args2kwargs(self, args, i=0):
@@ -284,6 +302,19 @@ def args2kwargs(self, args, i=0):
                 gamma_pl_list.append(args[i])
                 i += 1
             kwargs["gamma_pl_list"] = gamma_pl_list
+        if self._Ddt_sampling_num > 0:
+            Ddt_list = []
+            for k in range(self._Ddt_sampling_num):
+                Ddt_list.append(args[i])
+                i += 1
+            kwargs["Ddt_list"] = Ddt_list
+        if self._Dd_sampling_num > 0:
+            Dd_list = []
+            for k in range(self._Dd_sampling_num):
+                Dd_list.append(args[i])
+                i += 1
+            kwargs["Dd_list"] = Dd_list
+
         if self._gamma_pl_global_sampling is True:
             if "gamma_pl_mean" in self._kwargs_fixed:
                 kwargs["gamma_pl_mean"] = self._kwargs_fixed["gamma_pl_mean"]
@@ -357,6 +388,12 @@ def kwargs2args(self, kwargs):
         if self._gamma_pl_num > 0:
             for i in range(self._gamma_pl_num):
                 args.append(kwargs["gamma_pl_list"][i])
+        if self._Ddt_sampling_num > 0:
+            for i in range(self._Ddt_sampling_num):
+                args.append(kwargs["Ddt_list"][i])
+        if self._Dd_sampling_num > 0:
+            for i in range(self._Dd_sampling_num):
+                args.append(kwargs["Dd_list"][i])
         if self._gamma_pl_global_sampling is True:
             if "gamma_pl_mean" not in self._kwargs_fixed:
                 args.append(kwargs["gamma_pl_mean"])

hierarc/Sampling/ParamManager/param_manager.py

@@ -32,6 +32,8 @@ def __init__(
         gamma_pl_num=0,
         gamma_pl_global_sampling=False,
         gamma_pl_global_dist="NONE",
+        Ddt_sampling_num=0,
+        Dd_sampling_num=0,
         sigma_v_systematics=False,
         sne_apparent_m_sampling=False,
         sne_distribution="GAUSSIAN",
@@ -133,6 +135,8 @@ def __init__(
             gamma_pl_num=gamma_pl_num,
             gamma_pl_global_sampling=gamma_pl_global_sampling,
             gamma_pl_global_dist=gamma_pl_global_dist,
+            Ddt_sampling_num=Ddt_sampling_num,
+            Dd_sampling_num=Dd_sampling_num,
             log_scatter=log_scatter,
             kwargs_fixed=kwargs_fixed_lens,
         )