Merge remote-tracking branch 'origin/dsp_msd' into dsp_msd

hierarc/Likelihood/LensLikelihood/base_lens_likelihood.py

@@ -1,6 +1,8 @@
 __author__ = "sibirrer"
 
-from hierarc.Likelihood.LensLikelihood.double_source_plane import beta_double_source_plane
+from hierarc.Likelihood.LensLikelihood.double_source_plane import (
+    beta_double_source_plane,
+)
 
 LIKELIHOOD_TYPES = [
     "DdtGaussian",
@@ -102,7 +104,9 @@ def __init__(
                 DdtHistLikelihood,
             )
 
-            self._lens_type = DdtHistLikelihood(z_lens, z_source, normalized=normalized, **kwargs_likelihood)
+            self._lens_type = DdtHistLikelihood(
+                z_lens, z_source, normalized=normalized, **kwargs_likelihood
+            )
         elif likelihood_type == "DdtHistKDE":
             from hierarc.Likelihood.LensLikelihood.ddt_hist_likelihood import (
                 DdtHistKDELikelihood,
@@ -146,7 +150,10 @@ def __init__(
 
             self._lens_type = TDMagMagnitudeLikelihood(**kwargs_likelihood)
         elif likelihood_type == "DSPL":
-            from hierarc.Likelihood.LensLikelihood.double_source_plane import DSPLikelihood
+            from hierarc.Likelihood.LensLikelihood.double_source_plane import (
+                DSPLikelihood,
+            )
+
             self._lens_type = DSPLikelihood(normalized=normalized, **kwargs_likelihood)
         else:
             raise ValueError(
@@ -162,8 +169,15 @@ def num_data(self):
         return self._lens_type.num_data
 
     def log_likelihood(
-        self, ddt, dd, beta_dsp=None, kin_scaling=None, sigma_v_sys_error=None, mu_intrinsic=None,
-        gamma_pl=None, lambda_mst=None
+        self,
+        ddt,
+        dd,
+        beta_dsp=None,
+        kin_scaling=None,
+        sigma_v_sys_error=None,
+        mu_intrinsic=None,
+        gamma_pl=None,
+        lambda_mst=None,
     ):
         """
 
@@ -201,7 +215,9 @@ def log_likelihood(
         elif self.likelihood_type in ["TDMag", "TDMagMagnitude"]:
             return self._lens_type.log_likelihood(ddt=ddt, mu_intrinsic=mu_intrinsic)
         elif self.likelihood_type in ["DSPL"]:
-            return self._lens_type.log_likelihood(beta_dsp=beta_dsp, gamma_pl=gamma_pl, lambda_mst=lambda_mst)
+            return self._lens_type.log_likelihood(
+                beta_dsp=beta_dsp, gamma_pl=gamma_pl, lambda_mst=lambda_mst
+            )
         else:
             raise ValueError(
                 "likelihood type %s not fully supported." % self.likelihood_type
@@ -249,14 +265,18 @@ def sigma_v_prediction(self, ddt, dd, kin_scaling=None):
 
     def beta_dsp(self, cosmo):
         """Model prediction of ratio of Einstein radii theta_E_1 / theta_E_2 or scaled
-        deflection angles. Only computes it when likelihood is DSP
+        deflection angles. Only computes it when likelihood is DSP.
 
         :param cosmo: ~astropy.cosmology instance
         :return: beta
         """
         if self.likelihood_type == "DSPL":
-            beta = beta_double_source_plane(z_lens=self.z_lens, z_source_1=self.z_source, z_source_2=self.z_source2,
-                                            cosmo=cosmo)
+            beta = beta_double_source_plane(
+                z_lens=self.z_lens,
+                z_source_1=self.z_source,
+                z_source_2=self.z_source2,
+                cosmo=cosmo,
+            )
         else:
             beta = None
         return beta

hierarc/Likelihood/LensLikelihood/double_source_plane.py

@@ -28,13 +28,15 @@ def log_likelihood(
     ):
         """Log likelihood of the data given a model.
 
-        :param beta_dsp: scaled deflection angles alpha_1 / alpha_2 as ratio
-         between z_source and z_source2 source planes
+        :param beta_dsp: scaled deflection angles alpha_1 / alpha_2 as ratio between
+            z_source and z_source2 source planes
         :param gamma_pl: power-law density slope of main deflector (=2 being isothermal)
         :param lambda_mst: mass-sheet transform at the main deflector
         :return: log likelihood of data given model
         """
-        theta_E_ratio = beta2theta_e_ratio(beta_dsp, gamma_pl=gamma_pl, lambda_mst=lambda_mst)
+        theta_E_ratio = beta2theta_e_ratio(
+            beta_dsp, gamma_pl=gamma_pl, lambda_mst=lambda_mst
+        )
         log_l = -0.5 * ((theta_E_ratio - self._beta_dspl) / self._sigma_beta_dspl) ** 2
         if self._normalized:
             log_l -= 1 / 2.0 * np.log(2 * np.pi * self._sigma_beta_dspl**2)
@@ -67,13 +69,13 @@ def beta_double_source_plane(z_lens, z_source_1, z_source_2, cosmo):
 
 
 def beta2theta_e_ratio(beta_dsp, gamma_pl=2, lambda_mst=1):
-    """
-    calculates Einstein radii ratio for a power-law + MST profile with given parameters
+    """Calculates Einstein radii ratio for a power-law + MST profile with given
+    parameters.
 
-    :param beta_dsp: scaled deflection angles alpha_1 / alpha_2 as ratio
-     between z_source and z_source2 source planes
+    :param beta_dsp: scaled deflection angles alpha_1 / alpha_2 as ratio between
+        z_source and z_source2 source planes
     :param gamma_pl: power-law density slope of main deflector (=2 being isothermal)
     :param lambda_mst: mass-sheet transform at the main deflector
     :return: theta_E1 / theta_E2
     """
-    return (beta_dsp - (1 - lambda_mst) * (1 - beta_dsp)) ** (1/(gamma_pl - 1))
+    return (beta_dsp - (1 - lambda_mst) * (1 - beta_dsp)) ** (1 / (gamma_pl - 1))

hierarc/Likelihood/hierarchy_likelihood.py

@@ -223,19 +223,22 @@ def hyper_param_likelihood(
         kwargs_los=None,
         cosmo=None,
     ):
-        """Log likelihood of the data of a lens given a model (defined with hyperparameters) and cosmological distances.
+        """Log likelihood of the data of a lens given a model (defined with
+        hyperparameters) and cosmological distances.
 
         :param ddt: time-delay distance
         :param dd: angular diameter distance to the deflector
         :param delta_lum_dist: relative luminosity distance to pivot redshift
-        :param beta_dsp: Model prediction of ratio of Einstein radii theta_E_1 / theta_E_2
+        :param beta_dsp: Model prediction of ratio of Einstein radii theta_E_1 /
+            theta_E_2
         :param kwargs_lens: keywords of the hyperparameters of the lens model
         :param kwargs_kin: keyword arguments of the kinematic model hyperparameters
         :param kwargs_source: keyword argument of the source model (such as SNe)
         :param kwargs_los: list of keyword arguments of global line of sight
             distributions
         :param cosmo: ~astropy.cosmology instance
-        :return: log likelihood given the single lens analysis for the given hyperparameter
+        :return: log likelihood given the single lens analysis for the given
+            hyperparameter
         """
         kwargs_lens = self._kwargs_init(kwargs_lens)
         kwargs_kin = self._kwargs_init(kwargs_kin)
@@ -296,7 +299,8 @@ def log_likelihood_single(
         :param ddt: time-delay distance
         :param dd: angular diameter distance to the deflector
         :param delta_lum_dist: relative luminosity distance to pivot redshift
-        :param beta_dsp: Model prediction of ratio of Einstein radii theta_E_1 / theta_E_2
+        :param beta_dsp: Model prediction of ratio of Einstein radii theta_E_1 /
+            theta_E_2
         :param kwargs_lens: keywords of the hyperparameters of the lens model
         :param kwargs_kin: keyword arguments of the kinematic model hyperparameters
         :param kwargs_source: keyword arguments of source brightness

hierarc/Likelihood/lens_sample_likelihood.py

@@ -23,7 +23,9 @@ def __init__(self, kwargs_lens_list, normalized=False, kwargs_global_model=None)
         self._gamma_pl_num = 0
         gamma_pl_index = 0
 
-        gamma_pl_global_sampling = kwargs_global_model.get("gamma_pl_global_sampling", False)
+        gamma_pl_global_sampling = kwargs_global_model.get(
+            "gamma_pl_global_sampling", False
+        )
 
         for kwargs_lens in kwargs_lens_list:
             gamma_pl_index_ = None
@@ -37,7 +39,11 @@ def __init__(self, kwargs_lens_list, normalized=False, kwargs_global_model=None)
                 kwargs_global_model=kwargs_global_model, kwargs_lens=kwargs_lens
             )
             self._lens_list.append(
-                LensLikelihood(gamma_pl_index=gamma_pl_index_, normalized=normalized, **kwargs_lens_)
+                LensLikelihood(
+                    gamma_pl_index=gamma_pl_index_,
+                    normalized=normalized,
+                    **kwargs_lens_
+                )
             )
 
     def log_likelihood(

notebooks/double_source_plane.ipynb

@@ -55,7 +55,8 @@
    "outputs": [],
    "source": [
     "from hierarc.Likelihood.LensLikelihood.double_source_plane import (\n",
-    "    beta_double_source_plane, beta2theta_e_ratio\n",
+    "    beta_double_source_plane,\n",
+    "    beta2theta_e_ratio,\n",
     ")\n",
     "\n",
     "# define a cosmology\n",
@@ -94,8 +95,12 @@
     "    # TODO: translate to measured Einstein radius ratio\n",
     "    lambda_mst = np.random.normal(lambda_mst_mean, lambda_mst_sigma)\n",
     "    gamma_pl = np.random.normal(gamma_pl_mean, gamma_pl_sigma)\n",
-    "    beta_transformed = beta2theta_e_ratio(beta_dsp=beta, gamma_pl=gamma_pl, lambda_mst=lambda_mst)\n",
-    "    beta_measured = beta_transformed + np.random.normal(loc=0, scale=sigma_beta * beta_transformed)\n",
+    "    beta_transformed = beta2theta_e_ratio(\n",
+    "        beta_dsp=beta, gamma_pl=gamma_pl, lambda_mst=lambda_mst\n",
+    "    )\n",
+    "    beta_measured = beta_transformed + np.random.normal(\n",
+    "        loc=0, scale=sigma_beta * beta_transformed\n",
+    "    )\n",
     "    kwargs_likelihood = {\n",
     "        \"z_lens\": z_lens,\n",
     "        \"z_source\": z_source1,\n",
@@ -110,7 +115,9 @@
     "kwargs_dspl_list = []\n",
     "\n",
     "for i in range(num_dspl):\n",
-    "    kwargs_dspl_list.append(draw_lens(lambda_mst_mean, lambda_mst_sigma, gamma_pl_mean, gamma_pl_sigma))"
+    "    kwargs_dspl_list.append(\n",
+    "        draw_lens(lambda_mst_mean, lambda_mst_sigma, gamma_pl_mean, gamma_pl_sigma)\n",
+    "    )"
    ]
   },
   {
@@ -135,23 +142,37 @@
     "\n",
     "cosmology = \"FwCDM\"\n",
     "\n",
-    "kwargs_mean_start = {\"kwargs_cosmo\": {\"h0\": 70, \"om\": 0.3, \"w\": -1},\n",
-    "                    \"kwargs_lens\": {\"lambda_mst\": lambda_mst_mean, \"lambda_mst_sigma\": lambda_mst_sigma,\n",
-    "                                    \"gamma_pl_mean\": gamma_pl_mean, \"gamma_pl_sigma\": gamma_pl_sigma}}\n",
+    "kwargs_mean_start = {\n",
+    "    \"kwargs_cosmo\": {\"h0\": 70, \"om\": 0.3, \"w\": -1},\n",
+    "    \"kwargs_lens\": {\n",
+    "        \"lambda_mst\": lambda_mst_mean,\n",
+    "        \"lambda_mst_sigma\": lambda_mst_sigma,\n",
+    "        \"gamma_pl_mean\": gamma_pl_mean,\n",
+    "        \"gamma_pl_sigma\": gamma_pl_sigma,\n",
+    "    },\n",
+    "}\n",
     "\n",
-    "kwargs_sigma_start = {\"kwargs_cosmo\": {\"h0\": 10, \"om\": 0.05, \"w\": 0.2},\n",
-    "                     \"kwargs_lens\": {\"lambda_mst\": 0.1, \"lambda_mst_sigma\": 0.01,\n",
-    "                                     \"gamma_pl_mean\": 0.1, \"gamma_pl_sigma\": 0.01}}\n",
+    "kwargs_sigma_start = {\n",
+    "    \"kwargs_cosmo\": {\"h0\": 10, \"om\": 0.05, \"w\": 0.2},\n",
+    "    \"kwargs_lens\": {\n",
+    "        \"lambda_mst\": 0.1,\n",
+    "        \"lambda_mst_sigma\": 0.01,\n",
+    "        \"gamma_pl_mean\": 0.1,\n",
+    "        \"gamma_pl_sigma\": 0.01,\n",
+    "    },\n",
+    "}\n",
     "\n",
     "\n",
     "kwargs_bounds = {\n",
     "    \"kwargs_lower_cosmo\": {\"h0\": 0, \"om\": 0, \"w\": -2},\n",
     "    \"kwargs_upper_cosmo\": {\"h0\": 200, \"om\": 1, \"w\": 0},\n",
     "    \"kwargs_fixed_cosmo\": {\"h0\": kwargs_cosmo_true[\"h0\"]},\n",
-    "    \n",
     "    \"kwargs_lower_lens\": {\"lambda_mst\": 0.8, \"gamma_pl_mean\": 1.5},\n",
     "    \"kwargs_upper_lens\": {\"lambda_mst\": 1.2, \"gamma_pl_mean\": 2.5},\n",
-    "    \"kwargs_fixed_lens\": {\"gamma_pl_sigma\": gamma_pl_sigma, \"lambda_mst_sigma\": lambda_mst_sigma}\n",
+    "    \"kwargs_fixed_lens\": {\n",
+    "        \"gamma_pl_sigma\": gamma_pl_sigma,\n",
+    "        \"lambda_mst_sigma\": lambda_mst_sigma,\n",
+    "    },\n",
     "}\n",
     "\n",
     "\n",
@@ -161,7 +182,6 @@
     "# joint options for hierArc sampling\n",
     "\n",
     "\n",
-    "\n",
     "kwargs_model = {\n",
     "    \"lambda_mst_sampling\": True,\n",
     "    \"lambda_mst_distribution\": \"NONE\",  # TODO: switch to GAUSSIAN\n",
@@ -171,11 +191,12 @@
     "}\n",
     "\n",
     "\n",
-    "kwargs_sampler = {\"custom_prior\": None,\n",
-    "                  \"interpolate_cosmo\": False,\n",
-    "                  \"num_redshift_interp\": 100,\n",
-    "                  \"cosmo_fixed\": None,\n",
-    "                 }"
+    "kwargs_sampler = {\n",
+    "    \"custom_prior\": None,\n",
+    "    \"interpolate_cosmo\": False,\n",
+    "    \"num_redshift_interp\": 100,\n",
+    "    \"cosmo_fixed\": None,\n",
+    "}"
    ]
   },
   {
@@ -193,20 +214,30 @@
     }
    ],
    "source": [
-    "mcmc_sampler = MCMCSampler(kwargs_dspl_list, cosmology=cosmology, \n",
-    "                           kwargs_model=kwargs_model, kwargs_bounds=kwargs_bounds,\n",
-    "                          **kwargs_sampler)\n",
+    "mcmc_sampler = MCMCSampler(\n",
+    "    kwargs_dspl_list,\n",
+    "    cosmology=cosmology,\n",
+    "    kwargs_model=kwargs_model,\n",
+    "    kwargs_bounds=kwargs_bounds,\n",
+    "    **kwargs_sampler\n",
+    ")\n",
     "\n",
     "likelihood = mcmc_sampler.chain\n",
     "param = mcmc_sampler.param\n",
     "\n",
-    "kwargs_test = {\"kwargs_cosmo\": {\"h0\": 70, \"om\": 0.3, \"w\": -1},\n",
-    "                    \"kwargs_lens\": {\"lambda_mst\": lambda_mst_mean, \"lambda_mst_sigma\": lambda_mst_sigma,\n",
-    "                                    \"gamma_pl_mean\": gamma_pl_mean+0., \"gamma_pl_sigma\": gamma_pl_sigma}}\n",
+    "kwargs_test = {\n",
+    "    \"kwargs_cosmo\": {\"h0\": 70, \"om\": 0.3, \"w\": -1},\n",
+    "    \"kwargs_lens\": {\n",
+    "        \"lambda_mst\": lambda_mst_mean,\n",
+    "        \"lambda_mst_sigma\": lambda_mst_sigma,\n",
+    "        \"gamma_pl_mean\": gamma_pl_mean + 0.0,\n",
+    "        \"gamma_pl_sigma\": gamma_pl_sigma,\n",
+    "    },\n",
+    "}\n",
     "\n",
     "args = param.kwargs2args(**kwargs_test)\n",
     "\n",
-    "%timeit likelihood.likelihood(args)\n"
+    "%timeit likelihood.likelihood(args)"
    ]
   },
   {
@@ -262,8 +293,6 @@
     }
    ],
    "source": [
-    "\n",
-    "\n",
     "mcmc_samples, log_prob = mcmc_sampler.mcmc_emcee(\n",
     "    n_walkers, n_burn, n_run, kwargs_mean_start, kwargs_sigma_start\n",
     ")\n",

test/test_Sampling/test_ParamManager/test_lens_param.py

@@ -78,7 +78,7 @@ def test_args2kwargs(self):
             "alpha_lambda": 0.1,
             "beta_lambda": 0.1,
             "gamma_pl_list": [2, 2.5],
-            "gamma_pl_mean": 2.,
+            "gamma_pl_mean": 2.0,
             "gamma_pl_sigma": 0.1,
         }
         args = self._param.kwargs2args(kwargs)