adding the DESI BAO likelihood

hierarc/Data/BAO/desi_bao_dr2/README.rst

@@ -0,0 +1,5 @@
+# DESI DR2 data set
+
+Source: https://github.com/CobayaSampler/bao_data/tree/master/desi_bao_dr2
+
+Reference Paper : https://arxiv.org/abs/2503.14738

hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_cov.txt

@@ -0,0 +1,13 @@
+5.78998687e-03 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 2.83473742e-02 -3.26062007e-02 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 -3.26062007e-02 1.83928040e-01 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 3.23752442e-02 -2.37445646e-02 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 -2.37445646e-02 1.11469198e-01 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 2.61732816e-02 -1.12938006e-02 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 -1.12938006e-02 4.04183878e-02 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 1.05336516e-01 -2.90308418e-02 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 -2.90308418e-02 5.04233092e-02 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 5.83020277e-01 -1.95215562e-01 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 -1.95215562e-01 2.68336193e-01 0.00000000e+00 0.00000000e+00
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 1.02136194e-02 -2.31395216e-02
+0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00 -2.31395216e-02 2.82685779e-01

hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_mean.txt

@@ -0,0 +1,14 @@
+# [z] [value at z] [quantity]
+0.29500000 7.94167639 DV_over_rs
+0.51000000 13.58758434 DM_over_rs
+0.51000000 21.86294686 DH_over_rs
+0.70600000 17.35069094 DM_over_rs
+0.70600000 19.45534918 DH_over_rs
+0.93400000 21.57563956 DM_over_rs
+0.93400000 17.64149464 DH_over_rs
+1.32100000 27.60085612 DM_over_rs
+1.32100000 14.17602155 DH_over_rs
+1.48400000 30.51190063 DM_over_rs
+1.48400000 12.81699964 DH_over_rs
+2.33 8.631545674846294 DH_over_rs
+2.33 38.988973961958784 DM_over_rs

hierarc/Likelihood/BAOLikelihood/bao_likelihood.py

@@ -0,0 +1,45 @@
+import numpy as np
+from hierarc.Likelihood.BAOLikelihood.bao_likelihood_custom import CustomBAOLikelihood
+
+
+class BAOLikelihood(object):
+    """BAO likelihood.This class supports custom likelihoods as well as likelihoods from
+    the DESI BAO files."""
+
+    def __init__(self, sample_name="DESI_DR2", **kwargs_bao_likelihood):
+        """
+
+        :param sample_name: string, either 'CUSTOM' or a specific name supported by SneLikelihoodFromFile() class
+        :param kwargs_sne_likelihood: keyword arguments to initiate likelihood class
+        """
+        if sample_name == "CUSTOM":
+            self._likelihood = CustomBAOLikelihood(**kwargs_bao_likelihood)
+        elif sample_name == "DESI_DR2":
+            from hierarc.Likelihood.BAOLikelihood.desi_dr2 import DESIDR2Data
+
+            data = DESIDR2Data()
+
+            self._likelihood = CustomBAOLikelihood(
+                z=data.z,
+                d=data.d,
+                distance_type=data.distance_type,
+                cov=data.cov,
+            )
+
+        else:
+            raise ValueError("Unsupported sample name: {}".format(sample_name))
+
+    def log_likelihood(self, cosmo, rd=None):
+        """
+
+        :param cosmo: instance of a class to compute angular diameter distances on arrays
+
+        :return: log likelihood of the data given the specified cosmology
+        """
+        # TODO compute here the default case if rd is not sampled.
+        if rd is None:
+            raise NotImplementedError(
+                "Computation of rd is not implemented yet. Please provide rd in the kwargs_cosmo and turn rd_sampling=True in the kwargs_model."
+            )
+
+        return self._likelihood.log_likelihood_bao(cosmo, rd)

hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py

@@ -0,0 +1,79 @@
+import numpy as np
+from astropy.constants import c
+import astropy.units as u
+
+_twopi = 2 * np.pi
+
+
+class CustomBAOLikelihood(object):
+    """Class method for an arbitrary BAO measurements.
+
+    Distances measurements (scaled by rs) and the covariance matrix must be provided in
+    the constructor. The likelihood is assumed to be Gaussian.
+    """
+
+    def __init__(self, z, d, distance_type, cov):
+        """
+
+        :param z: array of redshifts of the BAO measurements
+        :param d: array of BAO measurements, scaled by rs
+        :param distance_type: string, either 'DV_over_rs' or 'DM_over_rs' or 'DH_over_rs'
+
+        """
+        self.z = z
+        self.d = d
+        self.distance_type = distance_type
+        self.cov = cov
+        self._inv_cov = np.linalg.inv(cov)
+        self.num_d = len(d)
+        assert len(z) == len(d), "z and d must have the same length"
+
+    def log_likelihood_bao(self, cosmo, rd):
+        """
+        :param cosmo: instance of a class to compute angular diameter distances on arrays
+        :param rd: comoving sound horizon at the drag epoch
+        :return: log likelihood of the data given the specified cosmology
+        """
+        distance_theory = np.zeros(self.num_d)
+
+        for i in range(self.num_d):
+            if self.distance_type[i] == "DV_over_rs":
+                distance_theory[i] = self._compute_DV(cosmo, self.z[i])
+            elif self.distance_type[i] == "DM_over_rs":
+                distance_theory[i] = self._compute_DM(cosmo, self.z[i])
+            elif self.distance_type[i] == "DH_over_rs":
+                distance_theory[i] = self._compute_DH(cosmo, self.z[i])
+            else:
+                raise ValueError(
+                    "Unsupported distance type: {}".format(self.distance_type)
+                )
+            # scale by the comoving sound horizon
+            distance_theory[i] /= rd
+
+        # Compute the log likelihood
+        diff = self.d - distance_theory
+        logL = -0.5 * np.dot(diff, np.dot(self._inv_cov, diff))
+        sign_det, lndet = np.linalg.slogdet(self.cov)
+        logL -= 1 / 2.0 * (self.num_d * np.log(2 * np.pi) + lndet)  # normalization term
+        return logL
+
+    def _compute_DV(self, cosmo, z):
+        """Compute the DV distance at redshift z.
+
+        (see Section III.A of https://arxiv.org/pdf/2503.14738)
+        """
+
+        DV = (z * self._compute_DM(cosmo, z) ** 2 * self._compute_DH(cosmo, z)) ** (
+            1.0 / 3.0
+        )
+        return DV
+
+    def _compute_DM(self, cosmo, z):
+        """Compute the DM distance (transverse comoving distance) at redshift z."""
+        return cosmo.comoving_transverse_distance(z).value
+
+    def _compute_DH(self, cosmo, z):
+        """Compute the DH (Hubble distance) distance at redshift z."""
+        Hz = cosmo.H(z)  # in km/s/Mpc
+        D_H = (c / Hz).to(u.Mpc)
+        return D_H.value

hierarc/Likelihood/BAOLikelihood/desi_dr2.py

@@ -0,0 +1,36 @@
+import os
+import pandas as pd
+import numpy as np
+import hierarc
+
+_PATH_2_DATA = os.path.join(os.path.dirname(hierarc.__file__), "Data", "BAO")
+
+
+class DESIDR2Data(object):
+    """
+    This class collect the data from teh DESI DR2 analysis presented in DESI collaboration et al. 2025 (https://arxiv.org/abs/2503.14738)
+
+    The data covariances that are stored in hierArc are originally from `DESI DR2`_.
+    https://github.com/CobayaSampler/bao_data/tree/master/desi_bao_dr2
+
+    If you make use of these products, please cite `DESI collaboration et al. 2025`_
+    """
+
+    def __init__(self):
+        self._data_file = os.path.join(
+            _PATH_2_DATA, "desi_bao_dr2", "desi_gaussian_bao_ALL_GCcomb_mean.txt"
+        )
+        self._cov_file = os.path.join(
+            _PATH_2_DATA, "desi_bao_dr2", "desi_gaussian_bao_ALL_GCcomb_cov.txt"
+        )
+
+        data = pd.read_csv(self._data_file, sep=r"\s+")
+
+        self.origlen = len(data)
+        print(f"Importing {self.origlen} distances from DESI DR2 data.")
+
+        self.z = data.iloc[:, 0].to_numpy()
+        self.d = data.iloc[:, 1].to_numpy()
+        self.distance_type = data.iloc[:, 2].to_numpy()
+
+        self.cov = np.loadtxt(self._cov_file)

hierarc/Likelihood/cosmo_likelihood.py

@@ -3,6 +3,7 @@
 from lenstronomy.Cosmo.cosmo_interp import CosmoInterp
 from hierarc.Likelihood.SneLikelihood.sne_likelihood import SneLikelihood
 from hierarc.Likelihood.KDELikelihood.kde_likelihood import KDELikelihood
+from hierarc.Likelihood.BAOLikelihood.bao_likelihood import BAOLikelihood
 from hierarc.Likelihood.KDELikelihood.chain import rescale_vector_to_unity
 import numpy as np
 
@@ -17,9 +18,11 @@ def __init__(
         kwargs_model,
         kwargs_bounds,
         sne_likelihood=None,
+        bao_likelihood=None,
         kwargs_sne_likelihood=None,
         KDE_likelihood_chain=None,
         kwargs_kde_likelihood=None,
+        kwargs_bao_likelihood=None,
         normalized=False,
         custom_prior=None,
         interpolate_cosmo=True,
@@ -42,7 +45,10 @@ def __init__(
         :param kwargs_kde_likelihood: keyword argument for the KDE likelihood, see KDELikelihood module for options
         :param sne_likelihood: (string), optional. Sampling supernovae relative expansion history likelihood, see
          SneLikelihood module for options
+         :param bao_likelihood: (string), optional. Sampling of BAO relative expansion history likelihood, see
+         BAOLikelihood module for options
         :param kwargs_sne_likelihood: keyword argument for the SNe likelihood, see SneLikelihood module for options
+        :param kwargs_bao_likelihood: keyword argument for the BAO likelihood, see BAOLikelihood module for options
         :param custom_prior: None or a definition that takes the keywords from the CosmoParam conventions and returns a
          log likelihood value (e.g. prior)
         :param interpolate_cosmo: bool, if True, uses interpolated comoving distance in the calculation for speed-up
@@ -79,11 +85,26 @@ def __init__(
             self._sne_likelihood = SneLikelihood(
                 sample_name=sne_likelihood, **kwargs_sne_likelihood
             )
-            z_max = np.max(self._sne_likelihood.zcmb)
+            z_max = max(z_max, np.max(self._sne_likelihood.zcmb))
             self._sne_evaluate = True
         else:
             self._sne_evaluate = False
 
+        if bao_likelihood is not None:
+            if interpolate_cosmo is True:
+                raise NotImplementedError(
+                    "BAO likelihood does not support interpolated cosmology yet. Please set interpolate_cosmo=False."
+                )
+            if kwargs_bao_likelihood is None:
+                kwargs_bao_likelihood = {}
+            self._bao_likelihood = BAOLikelihood(
+                sample_name=bao_likelihood, **kwargs_bao_likelihood
+            )
+            z_max = max(z_max, np.max(self._bao_likelihood._likelihood.z))
+            self._bao_evaluate = True
+        else:
+            self._bao_evaluate = False
+
         if KDE_likelihood_chain is not None:
             if kwargs_kde_likelihood is None:
                 kwargs_kde_likelihood = {}
@@ -174,6 +195,9 @@ def likelihood(self, args, kwargs_cosmo_interp=None, verbose=False):
                 cosmo_params, self._kde_likelihood.chain.rescale_dic, self._chain_params
             )
             log_l += self._kde_likelihood.kdelikelihood_samples(cosmo_params)[0]
+        if self._bao_evaluate is True:
+            rd = kwargs_cosmo.get("rd", None)
+            log_l += self._bao_likelihood.log_likelihood(cosmo=cosmo, rd=rd)
         if self._prior_add is True:
             log_l += self._custom_prior(
                 kwargs_cosmo, kwargs_lens, kwargs_kin, kwargs_source, kwargs_los

hierarc/Likelihood/hierarchy_likelihood.py

@@ -128,7 +128,9 @@ def __init__(
         :param gamma_pl_pivot: pivote power-law slope relative to which the Ddt scaling operates
         :param gamma_pl_ddt_slope: slope of Ddt / gamma_pl such that Ddt(gamma_pl) = Ddt_0 + slope * gamma_pl
         """
-        TransformedCosmography.__init__(self, gamma_pl_pivot=gamma_pl_pivot, gamma_pl_ddt_slope=gamma_pl_ddt_slope)
+        TransformedCosmography.__init__(
+            self, gamma_pl_pivot=gamma_pl_pivot, gamma_pl_ddt_slope=gamma_pl_ddt_slope
+        )
 
         KinScaling.__init__(
             self,
@@ -530,11 +532,15 @@ def sigma_v_measured_vs_predict(
             lambda_mst, gamma_ppn, gamma_pl = (
                 kwargs_lens_draw["lambda_mst"],
                 kwargs_lens_draw["gamma_ppn"],
-                kwargs_lens_draw.get("gamma_pl", None)
+                kwargs_lens_draw.get("gamma_pl", None),
             )
             kappa_ext = self._los.draw_los(kwargs_los, size=1)[0]
             ddt_, dd_, _ = self.displace_prediction(
-                ddt, dd, gamma_ppn=gamma_ppn, lambda_mst=lambda_mst, kappa_ext=kappa_ext,
+                ddt,
+                dd,
+                gamma_ppn=gamma_ppn,
+                lambda_mst=lambda_mst,
+                kappa_ext=kappa_ext,
                 gamma_pl=gamma_pl,
             )
             kwargs_kin_draw = self._aniso_distribution.draw_anisotropy(
@@ -589,7 +595,12 @@ def ddt_dd_model_prediction(self, cosmo, kwargs_lens=None, kwargs_los=None):
             )
             kappa_ext = self._los.draw_los(kwargs_los)
             ddt_, dd_, _ = self.displace_prediction(
-                ddt, dd, gamma_ppn=gamma_ppn, lambda_mst=lambda_mst, kappa_ext=kappa_ext, gamma_pl=gamma_pl
+                ddt,
+                dd,
+                gamma_ppn=gamma_ppn,
+                lambda_mst=lambda_mst,
+                kappa_ext=kappa_ext,
+                gamma_pl=gamma_pl,
             )
             ddt_draws.append(ddt_)
             dd_draws.append(dd_)

hierarc/Likelihood/transformed_cosmography.py

@@ -20,7 +20,14 @@ def __init__(self, gamma_pl_pivot=None, gamma_pl_ddt_slope=None):
         self._gamma_pl_ddt_slope = gamma_pl_ddt_slope
 
     def displace_prediction(
-        self, ddt, dd, gamma_ppn=1, lambda_mst=1, kappa_ext=0, mag_source=0, gamma_pl=None,
+        self,
+        ddt,
+        dd,
+        gamma_ppn=1,
+        lambda_mst=1,
+        kappa_ext=0,
+        mag_source=0,
+        gamma_pl=None,
     ):
         """Here we effectively change the posteriors of the lens, but rather than
         changing the instance of the KDE we displace the predicted angular diameter
@@ -116,9 +123,9 @@ def _displace_lambda_mst(ddt, dd, lambda_mst=1, kappa_ext=0, mag_source=0):
         return ddt_, dd_, mag_source_
 
     def _displace_gamma_pl(self, ddt, gamma_pl):
-        """
-        shifts Ddt prediction as a function of the power-law density slope.
-        This shift is the inverse such that the Ddt predictions in the likelihood are unchanged.
+        """Shifts Ddt prediction as a function of the power-law density slope. This
+        shift is the inverse such that the Ddt predictions in the likelihood are
+        unchanged.
 
         :param ddt: model-predicted time-delay distance
         :param gamma_pl: model-predicted power-law density slope