From 939c030127b670bf5b566e1e8760ca2d1f552811 Mon Sep 17 00:00:00 2001
From: martin-millon <martin.millon@epfl.ch>
Date: Wed, 21 May 2025 15:39:31 +0200
Subject: [PATCH 01/10] adding the DESI BAO likelihood

---
 .../desi_gaussian_bao_ALL_GCcomb_cov.txt      | 13 ++++
 .../desi_gaussian_bao_ALL_GCcomb_mean.txt     | 14 ++++
 hierarc/Likelihood/BAOLikelihood/__init__.py  |  0
 .../BAOLikelihood/bao_likelihood.py           | 43 +++++++++++
 .../BAOLikelihood/bao_likelihood_custom.py    | 71 +++++++++++++++++++
 hierarc/Likelihood/BAOLikelihood/desi_dr2.py  | 37 ++++++++++
 hierarc/Likelihood/cosmo_likelihood.py        | 19 +++++
 hierarc/Sampling/ParamManager/cosmo_param.py  | 22 +++++-
 .../Sampling/ParamManager/param_manager.py    |  3 +
 9 files changed, 221 insertions(+), 1 deletion(-)
 create mode 100644 hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_cov.txt
 create mode 100644 hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_mean.txt
 create mode 100644 hierarc/Likelihood/BAOLikelihood/__init__.py
 create mode 100644 hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
 create mode 100644 hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
 create mode 100644 hierarc/Likelihood/BAOLikelihood/desi_dr2.py

diff --git a/hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_cov.txt b/hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_cov.txt
new file mode 100644
index 00000000..fd8e5697
--- /dev/null
+++ b/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
diff --git a/hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_mean.txt b/hierarc/Data/BAO/desi_bao_dr2/desi_gaussian_bao_ALL_GCcomb_mean.txt
new file mode 100644
index 00000000..8aff444f
--- /dev/null
+++ b/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
diff --git a/hierarc/Likelihood/BAOLikelihood/__init__.py b/hierarc/Likelihood/BAOLikelihood/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py b/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
new file mode 100644
index 00000000..e1d42de4
--- /dev/null
+++ b/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
@@ -0,0 +1,43 @@
+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.")
+
+        return self._likelihood.log_likelihood_bao(cosmo, rd)
diff --git a/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
new file mode 100644
index 00000000..7571a291
--- /dev/null
+++ b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
@@ -0,0 +1,71 @@
+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
+        inv_cov = self._inv_cov
+        logL = -0.5 * np.dot(diff, np.dot(inv_cov, diff))
+        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./3.)
+        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
+
+
diff --git a/hierarc/Likelihood/BAOLikelihood/desi_dr2.py b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
new file mode 100644
index 00000000..be647f83
--- /dev/null
+++ b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
@@ -0,0 +1,37 @@
+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, delim_whitespace=True)
+
+        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)
+
diff --git a/hierarc/Likelihood/cosmo_likelihood.py b/hierarc/Likelihood/cosmo_likelihood.py
index c0f921b4..937ff733 100644
--- a/hierarc/Likelihood/cosmo_likelihood.py
+++ b/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
@@ -84,6 +90,16 @@ def __init__(
         else:
             self._sne_evaluate = False
 
+        if bao_likelihood is not None:
+            if kwargs_bao_likelihood is None:
+                kwargs_bao_likelihood = {}
+            self._bao_likelihood = BAOLikelihood(
+                sample_name=bao_likelihood, **kwargs_bao_likelihood
+            )
+            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 +190,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
diff --git a/hierarc/Sampling/ParamManager/cosmo_param.py b/hierarc/Sampling/ParamManager/cosmo_param.py
index 89460b2c..62998181 100644
--- a/hierarc/Sampling/ParamManager/cosmo_param.py
+++ b/hierarc/Sampling/ParamManager/cosmo_param.py
@@ -6,11 +6,12 @@
 class CosmoParam(object):
     """Manages the cosmological parameters in the sampling."""
 
-    def __init__(self, cosmology, ppn_sampling=False, kwargs_fixed=None):
+    def __init__(self, cosmology, ppn_sampling=False, rd_sampling=False, kwargs_fixed=None):
         """
 
         :param cosmology: string describing cosmological model
         :param ppn_sampling: post-newtonian parameter sampling
+        :param rd_sampling: sound horizon at drag epoch sampling (used only if the BAOLIkelihood is used)
         :param kwargs_fixed: keyword arguments of fixed parameters during sampling
         """
         self._cosmology = cosmology
@@ -18,6 +19,7 @@ def __init__(self, cosmology, ppn_sampling=False, kwargs_fixed=None):
             kwargs_fixed = {}
         self._kwargs_fixed = kwargs_fixed
         self._ppn_sampling = ppn_sampling
+        self._rd_sampling = rd_sampling
         self._supported_cosmologies = [
             "FLCDM",
             "FwCDM",
@@ -91,6 +93,14 @@ def param_list(self, latex_style=False):
                     list.append(r"$\gamma_{\rm ppn}$")
                 else:
                     list.append("gamma_ppn")
+
+        if self._rd_sampling is True:
+            if "rd" not in self._kwargs_fixed:
+                if latex_style is True:
+                    list.append(r"$r_d$")
+                else:
+                    list.append("rd")
+
         return list
 
     def args2kwargs(self, args, i=0):
@@ -152,6 +162,13 @@ def args2kwargs(self, args, i=0):
             else:
                 kwargs["gamma_ppn"] = args[i]
                 i += 1
+        if self._rd_sampling is True:
+            if "rd" in self._kwargs_fixed:
+                kwargs["rd"] = self._kwargs_fixed["rd"]
+            else:
+                kwargs["rd"] = args[i]
+                i += 1
+
         return kwargs, i
 
     def kwargs2args(self, kwargs):
@@ -186,6 +203,9 @@ def kwargs2args(self, kwargs):
         if self._ppn_sampling is True:
             if "gamma_ppn" not in self._kwargs_fixed:
                 args.append(kwargs["gamma_ppn"])
+        if self._rd_sampling is True:
+            if "rd" not in self._kwargs_fixed:
+                args.append(kwargs["rd"])
         return args
 
     def cosmo(self, kwargs):
diff --git a/hierarc/Sampling/ParamManager/param_manager.py b/hierarc/Sampling/ParamManager/param_manager.py
index 33d5c692..33c05b02 100644
--- a/hierarc/Sampling/ParamManager/param_manager.py
+++ b/hierarc/Sampling/ParamManager/param_manager.py
@@ -12,6 +12,7 @@ def __init__(
         self,
         cosmology,
         ppn_sampling=False,
+        rd_sampling=False,
         lambda_mst_sampling=False,
         lambda_mst_distribution="NONE",
         anisotropy_sampling=False,
@@ -58,6 +59,7 @@ def __init__(
 
         :param cosmology: string describing cosmological model
         :param ppn_sampling: post-newtonian parameter sampling
+        :param rd_sampling: sound horizon at drag epoch sampling (used only if the BAOLIkelihood is used)
         :param lambda_mst_sampling: bool, if True adds a global mass-sheet transform parameter in the sampling
         :param lambda_mst_distribution: string, distribution function of the MST transform
         :param lambda_ifu_sampling: bool, if True samples a separate lambda_mst for a second (e.g. IFU) data set
@@ -112,6 +114,7 @@ def __init__(
         self._cosmo_param = CosmoParam(
             cosmology=cosmology,
             ppn_sampling=ppn_sampling,
+            rd_sampling=rd_sampling,
             kwargs_fixed=kwargs_fixed_cosmo,
         )
         self._lens_param = LensParam(

From 91653d1d876744e82a1b625ef4c5ed28a23a1669 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 May 2025 13:45:27 +0000
Subject: [PATCH 02/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 .../BAOLikelihood/bao_likelihood.py           | 14 ++++++-----
 .../BAOLikelihood/bao_likelihood_custom.py    | 25 ++++++++++++-------
 hierarc/Likelihood/BAOLikelihood/desi_dr2.py  |  1 -
 hierarc/Likelihood/hierarchy_likelihood.py    | 19 +++++++++++---
 hierarc/Likelihood/transformed_cosmography.py | 15 ++++++++---
 hierarc/Sampling/ParamManager/cosmo_param.py  |  4 ++-
 6 files changed, 53 insertions(+), 25 deletions(-)

diff --git a/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py b/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
index e1d42de4..154d2bd4 100644
--- a/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
+++ b/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
@@ -1,9 +1,10 @@
 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."""
+    """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):
         """
@@ -15,6 +16,7 @@ def __init__(self, sample_name="DESI_DR2", **kwargs_bao_likelihood):
             self._likelihood = CustomBAOLikelihood(**kwargs_bao_likelihood)
         elif sample_name == "DESI_DR2":
             from hierarc.Likelihood.BAOLikelihood.desi_dr2 import DESIDR2Data
+
             data = DESIDR2Data()
 
             self._likelihood = CustomBAOLikelihood(
@@ -27,8 +29,6 @@ def __init__(self, sample_name="DESI_DR2", **kwargs_bao_likelihood):
         else:
             raise ValueError("Unsupported sample name: {}".format(sample_name))
 
-
-
     def log_likelihood(self, cosmo, rd=None):
         """
 
@@ -36,8 +36,10 @@ def log_likelihood(self, cosmo, rd=None):
 
         :return: log likelihood of the data given the specified cosmology
         """
-        #TODO compute here the default case if rd is not sampled.
+        # 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.")
+            raise NotImplementedError(
+                "Computation of rd is not implemented yet. Please provide rd."
+            )
 
         return self._likelihood.log_likelihood_bao(cosmo, rd)
diff --git a/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
index 7571a291..89fe4f1e 100644
--- a/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
+++ b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
@@ -6,9 +6,11 @@
 
 
 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."""
+    """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):
         """
@@ -42,8 +44,10 @@ def log_likelihood_bao(self, cosmo, rd):
             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
+                raise ValueError(
+                    "Unsupported distance type: {}".format(self.distance_type)
+                )
+            # scale by the comoving sound horizon
             distance_theory[i] /= rd
 
         # Compute the log likelihood
@@ -53,9 +57,14 @@ def log_likelihood_bao(self, cosmo, rd):
         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)"""
+        """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./3.)
+        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):
@@ -67,5 +76,3 @@ def _compute_DH(self, cosmo, z):
         Hz = cosmo.H(z)  # in km/s/Mpc
         D_H = (c / Hz).to(u.Mpc)
         return D_H.value
-
-
diff --git a/hierarc/Likelihood/BAOLikelihood/desi_dr2.py b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
index be647f83..c8e44474 100644
--- a/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
+++ b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
@@ -34,4 +34,3 @@ def __init__(self):
         self.distance_type = data.iloc[:, 2].to_numpy()
 
         self.cov = np.loadtxt(self._cov_file)
-
diff --git a/hierarc/Likelihood/hierarchy_likelihood.py b/hierarc/Likelihood/hierarchy_likelihood.py
index 2a8e907f..c0847291 100644
--- a/hierarc/Likelihood/hierarchy_likelihood.py
+++ b/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_)
diff --git a/hierarc/Likelihood/transformed_cosmography.py b/hierarc/Likelihood/transformed_cosmography.py
index 647b0f75..131f6070 100644
--- a/hierarc/Likelihood/transformed_cosmography.py
+++ b/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
diff --git a/hierarc/Sampling/ParamManager/cosmo_param.py b/hierarc/Sampling/ParamManager/cosmo_param.py
index 62998181..201e89ad 100644
--- a/hierarc/Sampling/ParamManager/cosmo_param.py
+++ b/hierarc/Sampling/ParamManager/cosmo_param.py
@@ -6,7 +6,9 @@
 class CosmoParam(object):
     """Manages the cosmological parameters in the sampling."""
 
-    def __init__(self, cosmology, ppn_sampling=False, rd_sampling=False, kwargs_fixed=None):
+    def __init__(
+        self, cosmology, ppn_sampling=False, rd_sampling=False, kwargs_fixed=None
+    ):
         """
 
         :param cosmology: string describing cosmological model

From 7028923913c10855bb45c159c02104c45d479e78 Mon Sep 17 00:00:00 2001
From: martin-millon <martin.millon@epfl.ch>
Date: Thu, 22 May 2025 11:00:43 +0200
Subject: [PATCH 03/10] adding tests + normalisation correction

---
 hierarc/Data/BAO/desi_bao_dr2/README.rst      |  5 ++
 .../BAOLikelihood/bao_likelihood_custom.py    |  5 +-
 .../test_BAOLikelihood/__init__.py            |  0
 .../test_BAOLikelihood/test_DESIDR2_data.py   | 22 +++++
 .../test_BAOLikelihood/test_bao_likelihood.py | 87 +++++++++++++++++++
 5 files changed, 117 insertions(+), 2 deletions(-)
 create mode 100644 hierarc/Data/BAO/desi_bao_dr2/README.rst
 create mode 100644 test/test_Likelihood/test_BAOLikelihood/__init__.py
 create mode 100644 test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py
 create mode 100644 test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py

diff --git a/hierarc/Data/BAO/desi_bao_dr2/README.rst b/hierarc/Data/BAO/desi_bao_dr2/README.rst
new file mode 100644
index 00000000..68970f38
--- /dev/null
+++ b/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
diff --git a/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
index 7571a291..18d4fed4 100644
--- a/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
+++ b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
@@ -48,8 +48,9 @@ def log_likelihood_bao(self, cosmo, rd):
 
         # Compute the log likelihood
         diff = self.d - distance_theory
-        inv_cov = self._inv_cov
-        logL = -0.5 * np.dot(diff, np.dot(inv_cov, diff))
+        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):
diff --git a/test/test_Likelihood/test_BAOLikelihood/__init__.py b/test/test_Likelihood/test_BAOLikelihood/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py b/test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py
new file mode 100644
index 00000000..0b84c97f
--- /dev/null
+++ b/test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py
@@ -0,0 +1,22 @@
+import numpy as np
+import numpy.testing as npt
+
+from hierarc.Likelihood.BAOLikelihood.desi_dr2 import DESIDR2Data
+
+class TestDESIDR2Data(object):
+    def setup_method(self):
+        pass
+
+    def test_import(self):
+        data = DESIDR2Data()
+        z = data.z
+        d = data.d
+        distance_type = data.distance_type
+        cov = data.cov
+
+        npt.assert_almost_equal(z[0], 0.29500)
+        assert len(z) == len(d)
+        assert len(z) == len(distance_type)
+        nx, ny = np.shape(cov)
+        assert nx == ny
+        assert nx == len(d)
diff --git a/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
new file mode 100644
index 00000000..742e1ed0
--- /dev/null
+++ b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
@@ -0,0 +1,87 @@
+from hierarc.Likelihood.BAOLikelihood.bao_likelihood import BAOLikelihood
+import pytest
+import numpy as np
+import numpy.testing as npt
+
+
+class TestBAO(object):
+    def setup_method(self):
+        np.random.seed(42)
+        # define redshifts
+        num = 8  # number of BAO measurements
+        z = np.linspace(start=0.1, stop=0.8, num=num)
+
+        # define cosmology
+        from astropy.cosmology import FlatLambdaCDM
+
+        om_mean, om_sigma = 0.3, 0.01
+        rd = 150 # comoving sound horizon at the drag epoch
+        cosmo_true = FlatLambdaCDM(H0=70, Om0=om_mean)
+
+        # compute BAO distances (DM only)
+        dist_true = cosmo_true.comoving_transverse_distance(z).value / rd
+        dist_type = ["DM_over_rs"] * num
+
+        # draw from scatter
+        sigma_d = np.sqrt(4e-02)
+        cov = np.diag(np.ones(num) * sigma_d**2)
+
+        dist_measured = np.random.multivariate_normal(dist_true, cov)
+        kwargs_bao_likelihood = {
+            "z": z,
+            "d": dist_measured,
+            "distance_type": dist_type,
+            "cov":cov,
+        }
+
+        self.likelihood = BAOLikelihood(sample_name="CUSTOM", **kwargs_bao_likelihood)
+        self.dists_true = dist_true
+        self.rd_true = rd
+        self.sigma_d_true = sigma_d
+        self.cosmo_true = cosmo_true
+
+    def test_log_likelihood(self):
+        logL = self.likelihood.log_likelihood(
+            self.cosmo_true,
+            self.rd_true,
+        )
+
+        logL_high = self.likelihood.log_likelihood(
+            self.cosmo_true,
+            self.rd_true + 10,
+        )
+        assert logL > logL_high
+
+        logL_low = self.likelihood.log_likelihood(
+            self.cosmo_true,
+            self.rd_true - 10,
+        )
+        print(logL_low, logL, logL_high)
+        assert logL > logL_low
+
+    def test_desi_dr2(self):
+        likelihood = BAOLikelihood(sample_name="DESI_DR2")
+        om_mean, om_sigma = 0.2975, 0.0086  # from DESI collaboration et al. 2025
+        from astropy.cosmology import FlatLambdaCDM
+
+        cosmo_mean = FlatLambdaCDM(H0=67.5, Om0=om_mean)
+        logL_mean = likelihood.log_likelihood(cosmo_mean, self.rd_true)
+
+        cosmo_sigma_plus = FlatLambdaCDM(H0=67.5, Om0=om_mean + om_sigma)
+        logL_sigma_plus = likelihood.log_likelihood(cosmo_sigma_plus, self.rd_true)
+
+        npt.assert_almost_equal(logL_sigma_plus - logL_mean, -0.8915, decimal=1)
+        cosmo_sigma_neg = FlatLambdaCDM(H0=67.5, Om0=om_mean - om_sigma)
+        logL_sigma_neg = likelihood.log_likelihood(cosmo_sigma_neg, self.rd_true)
+
+        npt.assert_almost_equal(logL_sigma_neg - logL_mean, -6.03403, decimal=1)
+
+    def test_raise(self):
+        with pytest.raises(ValueError):
+            BAOLikelihood(sample_name="UNKNOWN")
+
+        with pytest.raises(NotImplementedError):
+            self.likelihood.log_likelihood(self.cosmo_true)
+
+if __name__ == "__main__":
+    pytest.main()

From 9cfbaae06ad04d245435c20f2447011eac8533db Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 22 May 2025 09:01:45 +0000
Subject: [PATCH 04/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py    | 2 +-
 test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py | 1 +
 .../test_BAOLikelihood/test_bao_likelihood.py                | 5 +++--
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
index 43ee6b83..0b7b4553 100644
--- a/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
+++ b/hierarc/Likelihood/BAOLikelihood/bao_likelihood_custom.py
@@ -54,7 +54,7 @@ def log_likelihood_bao(self, cosmo, rd):
         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
+        logL -= 1 / 2.0 * (self.num_d * np.log(2 * np.pi) + lndet)  # normalization term
         return logL
 
     def _compute_DV(self, cosmo, z):
diff --git a/test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py b/test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py
index 0b84c97f..30529a5e 100644
--- a/test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py
+++ b/test/test_Likelihood/test_BAOLikelihood/test_DESIDR2_data.py
@@ -3,6 +3,7 @@
 
 from hierarc.Likelihood.BAOLikelihood.desi_dr2 import DESIDR2Data
 
+
 class TestDESIDR2Data(object):
     def setup_method(self):
         pass
diff --git a/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
index 742e1ed0..773dd7e4 100644
--- a/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
+++ b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
@@ -15,7 +15,7 @@ def setup_method(self):
         from astropy.cosmology import FlatLambdaCDM
 
         om_mean, om_sigma = 0.3, 0.01
-        rd = 150 # comoving sound horizon at the drag epoch
+        rd = 150  # comoving sound horizon at the drag epoch
         cosmo_true = FlatLambdaCDM(H0=70, Om0=om_mean)
 
         # compute BAO distances (DM only)
@@ -31,7 +31,7 @@ def setup_method(self):
             "z": z,
             "d": dist_measured,
             "distance_type": dist_type,
-            "cov":cov,
+            "cov": cov,
         }
 
         self.likelihood = BAOLikelihood(sample_name="CUSTOM", **kwargs_bao_likelihood)
@@ -83,5 +83,6 @@ def test_raise(self):
         with pytest.raises(NotImplementedError):
             self.likelihood.log_likelihood(self.cosmo_true)
 
+
 if __name__ == "__main__":
     pytest.main()

From 90cfe797c9893b16b6349354bb09d4f94f70b998 Mon Sep 17 00:00:00 2001
From: martin-millon <martin.millon@epfl.ch>
Date: Thu, 22 May 2025 13:11:57 +0200
Subject: [PATCH 05/10] coverage improvement

---
 hierarc/Likelihood/BAOLikelihood/desi_dr2.py  |  2 +-
 hierarc/Likelihood/cosmo_likelihood.py        |  3 ++-
 test/test_Likelihood/test_cosmo_likelihood.py | 20 ++++++++++++++++++-
 .../test_ParamManager/test_cosmo_param.py     |  4 ++++
 4 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/hierarc/Likelihood/BAOLikelihood/desi_dr2.py b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
index c8e44474..ec3e188f 100644
--- a/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
+++ b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
@@ -24,7 +24,7 @@ def __init__(self):
             _PATH_2_DATA, "desi_bao_dr2", "desi_gaussian_bao_ALL_GCcomb_cov.txt"
         )
 
-        data = pd.read_csv(self._data_file, delim_whitespace=True)
+        data = pd.read_csv(self._data_file, sep=r'\s+')
 
         self.origlen = len(data)
         print(f"Importing {self.origlen} distances from DESI DR2 data.")
diff --git a/hierarc/Likelihood/cosmo_likelihood.py b/hierarc/Likelihood/cosmo_likelihood.py
index 937ff733..e96e8cc0 100644
--- a/hierarc/Likelihood/cosmo_likelihood.py
+++ b/hierarc/Likelihood/cosmo_likelihood.py
@@ -85,7 +85,7 @@ 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
@@ -96,6 +96,7 @@ def __init__(
             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
diff --git a/test/test_Likelihood/test_cosmo_likelihood.py b/test/test_Likelihood/test_cosmo_likelihood.py
index c21b416f..dcbd1531 100644
--- a/test/test_Likelihood/test_cosmo_likelihood.py
+++ b/test/test_Likelihood/test_cosmo_likelihood.py
@@ -47,8 +47,9 @@ def setup_method(self):
             "w": -2,
             "wa": -1,
             "w0": -2,
+            "rd":0
         }
-        kwargs_upper_cosmo = {"h0": 200, "om": 1, "ok": 0.8, "w": 0, "wa": 1, "w0": 1}
+        kwargs_upper_cosmo = {"h0": 200, "om": 1, "ok": 0.8, "w": 0, "wa": 1, "w0": 1., "rd":300}
         self.cosmology = "oLCDM"
         self.kwargs_bounds = {
             "kwargs_lower_lens": kwargs_lower_lens,
@@ -235,6 +236,23 @@ def test_sne_likelihood_integration(self):
         logl = cosmoL.likelihood(args=args, verbose=True)
         assert logl < 0
 
+    def test_bao_likelihood_integration(self):
+        kwargs_model = self.kwargs_model.copy()
+        kwargs_model["rd_sampling"] = True
+        cosmoL = CosmoLikelihood(
+            [],
+            self.cosmology,
+            kwargs_model,
+            self.kwargs_bounds,
+            bao_likelihood="DESI_DR2",
+            interpolate_cosmo=False,
+            cosmo_fixed=None,
+        )
+        kwargs_cosmo = {"h0": self.H0_true, "om": self.omega_m_true, "ok": 0, "rd":150}
+        args = cosmoL.param.kwargs2args(kwargs_cosmo=kwargs_cosmo)
+        logl = cosmoL.likelihood(args=args, verbose=True)
+        assert logl < 0
+
     def test_kde_likelihood_integration(self):
         chain = import_Planck_chain(
             _PATH_2_PLANCKDATA,
diff --git a/test/test_Sampling/test_ParamManager/test_cosmo_param.py b/test/test_Sampling/test_ParamManager/test_cosmo_param.py
index ca3b45b9..eecd3b62 100644
--- a/test/test_Sampling/test_ParamManager/test_cosmo_param.py
+++ b/test/test_Sampling/test_ParamManager/test_cosmo_param.py
@@ -16,6 +16,7 @@ def setup_method(self):
             "w0": -0,
             "gamma_ppn": 1,
             "alpha": 1.45,
+            "rd": 150.0,
         }
         self._param_list = []
         self._param_list_fixed = []
@@ -58,6 +59,7 @@ def test_args2kwargs(self):
             "w0": -0,
             "gamma_ppn": 1,
             "alpha": 1.45,
+            "rd": 150.0,
         }
         for i, param in enumerate(self._param_list):
             args = param.kwargs2args(kwargs)
@@ -80,6 +82,7 @@ def test_cosmo(self):
             "w0": -0,
             "gamma_ppn": 1,
             "alpha": 1.45,
+            "rd": 150.0,
         }
         for i, param in enumerate(self._param_list):
             cosmo = param.cosmo(kwargs_cosmo)
@@ -101,6 +104,7 @@ def test_raise(self):
                 "w0": -0,
                 "gamma_ppn": 1,
                 "alpha": 1.45,
+                "rd": 150.0,
             }
             param.cosmo(kwargs_cosmo)
 

From 3eb8d7fa6e65c870c1cd2106c523d7c0f7296569 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 22 May 2025 11:13:07 +0000
Subject: [PATCH 06/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 hierarc/Likelihood/BAOLikelihood/desi_dr2.py  |  2 +-
 hierarc/Likelihood/cosmo_likelihood.py        |  2 +-
 test/test_Likelihood/test_cosmo_likelihood.py | 14 +++++++++++---
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/hierarc/Likelihood/BAOLikelihood/desi_dr2.py b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
index ec3e188f..fbc65880 100644
--- a/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
+++ b/hierarc/Likelihood/BAOLikelihood/desi_dr2.py
@@ -24,7 +24,7 @@ def __init__(self):
             _PATH_2_DATA, "desi_bao_dr2", "desi_gaussian_bao_ALL_GCcomb_cov.txt"
         )
 
-        data = pd.read_csv(self._data_file, sep=r'\s+')
+        data = pd.read_csv(self._data_file, sep=r"\s+")
 
         self.origlen = len(data)
         print(f"Importing {self.origlen} distances from DESI DR2 data.")
diff --git a/hierarc/Likelihood/cosmo_likelihood.py b/hierarc/Likelihood/cosmo_likelihood.py
index e96e8cc0..d215b1ce 100644
--- a/hierarc/Likelihood/cosmo_likelihood.py
+++ b/hierarc/Likelihood/cosmo_likelihood.py
@@ -96,7 +96,7 @@ def __init__(
             self._bao_likelihood = BAOLikelihood(
                 sample_name=bao_likelihood, **kwargs_bao_likelihood
             )
-            z_max = max(z_max,np.max(self._bao_likelihood._likelihood.z))
+            z_max = max(z_max, np.max(self._bao_likelihood._likelihood.z))
             self._bao_evaluate = True
         else:
             self._bao_evaluate = False
diff --git a/test/test_Likelihood/test_cosmo_likelihood.py b/test/test_Likelihood/test_cosmo_likelihood.py
index dcbd1531..a6abaf36 100644
--- a/test/test_Likelihood/test_cosmo_likelihood.py
+++ b/test/test_Likelihood/test_cosmo_likelihood.py
@@ -47,9 +47,17 @@ def setup_method(self):
             "w": -2,
             "wa": -1,
             "w0": -2,
-            "rd":0
+            "rd": 0,
+        }
+        kwargs_upper_cosmo = {
+            "h0": 200,
+            "om": 1,
+            "ok": 0.8,
+            "w": 0,
+            "wa": 1,
+            "w0": 1.0,
+            "rd": 300,
         }
-        kwargs_upper_cosmo = {"h0": 200, "om": 1, "ok": 0.8, "w": 0, "wa": 1, "w0": 1., "rd":300}
         self.cosmology = "oLCDM"
         self.kwargs_bounds = {
             "kwargs_lower_lens": kwargs_lower_lens,
@@ -248,7 +256,7 @@ def test_bao_likelihood_integration(self):
             interpolate_cosmo=False,
             cosmo_fixed=None,
         )
-        kwargs_cosmo = {"h0": self.H0_true, "om": self.omega_m_true, "ok": 0, "rd":150}
+        kwargs_cosmo = {"h0": self.H0_true, "om": self.omega_m_true, "ok": 0, "rd": 150}
         args = cosmoL.param.kwargs2args(kwargs_cosmo=kwargs_cosmo)
         logl = cosmoL.likelihood(args=args, verbose=True)
         assert logl < 0

From 279858697635b9ff7712a9c7bbb6326a17d8041a Mon Sep 17 00:00:00 2001
From: martin-millon <martin.millon@epfl.ch>
Date: Thu, 22 May 2025 15:06:17 +0200
Subject: [PATCH 07/10] coverage improvement

---
 .../test_BAOLikelihood/test_bao_likelihood.py | 19 +++++++++++++------
 .../test_ParamManager/test_param_manager.py   | 10 ++++++++--
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
index 773dd7e4..205797da 100644
--- a/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
+++ b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
@@ -8,8 +8,8 @@ class TestBAO(object):
     def setup_method(self):
         np.random.seed(42)
         # define redshifts
-        num = 8  # number of BAO measurements
-        z = np.linspace(start=0.1, stop=0.8, num=num)
+        self.num = 8  # number of BAO measurements
+        z = np.linspace(start=0.1, stop=0.8, num=self.num)
 
         # define cosmology
         from astropy.cosmology import FlatLambdaCDM
@@ -20,21 +20,21 @@ def setup_method(self):
 
         # compute BAO distances (DM only)
         dist_true = cosmo_true.comoving_transverse_distance(z).value / rd
-        dist_type = ["DM_over_rs"] * num
+        dist_type = ["DM_over_rs"] * self.num
 
         # draw from scatter
         sigma_d = np.sqrt(4e-02)
-        cov = np.diag(np.ones(num) * sigma_d**2)
+        cov = np.diag(np.ones(self.num) * sigma_d**2)
 
         dist_measured = np.random.multivariate_normal(dist_true, cov)
-        kwargs_bao_likelihood = {
+        self.kwargs_bao_likelihood = {
             "z": z,
             "d": dist_measured,
             "distance_type": dist_type,
             "cov": cov,
         }
 
-        self.likelihood = BAOLikelihood(sample_name="CUSTOM", **kwargs_bao_likelihood)
+        self.likelihood = BAOLikelihood(sample_name="CUSTOM", **self.kwargs_bao_likelihood)
         self.dists_true = dist_true
         self.rd_true = rd
         self.sigma_d_true = sigma_d
@@ -83,6 +83,13 @@ def test_raise(self):
         with pytest.raises(NotImplementedError):
             self.likelihood.log_likelihood(self.cosmo_true)
 
+        kwargs_bao_likelihood = self.kwargs_bao_likelihood.copy()
+        kwargs_bao_likelihood["distance_type"] = ['Unknown'] * self.num
+        test_L = BAOLikelihood(sample_name="CUSTOM", **kwargs_bao_likelihood)
+
+        with pytest.raises(ValueError):
+            test_L.log_likelihood(self.cosmo_true, self.rd_true)
+
 
 if __name__ == "__main__":
     pytest.main()
diff --git a/test/test_Sampling/test_ParamManager/test_param_manager.py b/test/test_Sampling/test_ParamManager/test_param_manager.py
index a8967af9..01e1c8de 100644
--- a/test/test_Sampling/test_ParamManager/test_param_manager.py
+++ b/test/test_Sampling/test_ParamManager/test_param_manager.py
@@ -15,6 +15,7 @@ def setup_method(self):
             "wa": -1,
             "w0": -2,
             "gamma_ppn": 0,
+            "rd": 0.0,
         }
         kwargs_lower_lens = {
             "lambda_mst": 0,
@@ -32,6 +33,7 @@ def setup_method(self):
             "wa": 1,
             "w0": 1,
             "gamma_ppn": 5,
+            "rd": 300.0,
         }
         kwargs_upper_lens = {
             "lambda_mst": 2,
@@ -49,6 +51,7 @@ def setup_method(self):
             "wa": -0,
             "w0": -0,
             "gamma_ppn": 1,
+            "rd": 150.0,
         }
         kwargs_fixed_lens = {
             "lambda_mst": 1,
@@ -65,6 +68,7 @@ def setup_method(self):
                     cosmology=cosmology,
                     ppn_sampling=True,
                     lambda_mst_sampling=True,
+                    rd_sampling=True,
                     lambda_mst_distribution="GAUSSIAN",
                     anisotropy_distribution="GAUSSIAN",
                     anisotropy_sampling=True,
@@ -96,6 +100,7 @@ def setup_method(self):
                     cosmology=cosmology,
                     ppn_sampling=True,
                     lambda_mst_sampling=True,
+                    rd_sampling=True,
                     lambda_mst_distribution="GAUSSIAN",
                     anisotropy_distribution="GAUSSIAN",
                     anisotropy_sampling=True,
@@ -127,7 +132,7 @@ def test_num_param(self):
         list = self.param_list[0].param_list(latex_style=False)
         assert len(list) == 0
         num = self.param_list[1].num_param
-        assert num == 11
+        assert num == 12
         for param in self.param_list:
             list = param.param_list(latex_style=True)
             list = param.param_list(latex_style=False)
@@ -141,6 +146,7 @@ def test_kwargs2args(self):
             "wa": -0,
             "w0": -0,
             "gamma_ppn": 1,
+            "rd": 150.0,
         }
         kwargs_lens = {
             "lambda_mst": 1,
@@ -192,7 +198,7 @@ def test_param_bounds(self):
         assert len(lower_limit) == 0
         lower_limit, upper_limit = self.param_list[1].param_bounds
         print(self.param_list[1].param_list())
-        assert len(lower_limit) == 11
+        assert len(lower_limit) == 12
 
 
 class TestRaise(unittest.TestCase):

From d4e15df535c4da4160b431582cc7626c95603128 Mon Sep 17 00:00:00 2001
From: martin-millon <martin.millon@epfl.ch>
Date: Thu, 22 May 2025 15:07:16 +0200
Subject: [PATCH 08/10] coverage improvement

---
 test/test_Sampling/test_ParamManager/test_param_manager.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/test_Sampling/test_ParamManager/test_param_manager.py b/test/test_Sampling/test_ParamManager/test_param_manager.py
index 01e1c8de..949c1f67 100644
--- a/test/test_Sampling/test_ParamManager/test_param_manager.py
+++ b/test/test_Sampling/test_ParamManager/test_param_manager.py
@@ -188,6 +188,7 @@ def test_cosmo(self):
             "wa": -0,
             "w0": -0,
             "gamma_ppn": 1,
+            "rd": 150.0,
         }
         for param in self.param_list:
             cosmo = param.cosmo(kwargs_cosmo)

From 0ad7029e776aad8ce0736d1194c27b2da5b3fb3d Mon Sep 17 00:00:00 2001
From: martin-millon <martin.millon@epfl.ch>
Date: Thu, 22 May 2025 15:11:31 +0200
Subject: [PATCH 09/10] completing coverage

---
 hierarc/Likelihood/BAOLikelihood/bao_likelihood.py |  2 +-
 hierarc/Likelihood/cosmo_likelihood.py             |  4 ++++
 test/test_Likelihood/test_cosmo_likelihood.py      | 11 +++++++++++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py b/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
index 154d2bd4..b46ea3f1 100644
--- a/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
+++ b/hierarc/Likelihood/BAOLikelihood/bao_likelihood.py
@@ -39,7 +39,7 @@ def log_likelihood(self, cosmo, rd=None):
         # 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."
+                "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)
diff --git a/hierarc/Likelihood/cosmo_likelihood.py b/hierarc/Likelihood/cosmo_likelihood.py
index d215b1ce..74471026 100644
--- a/hierarc/Likelihood/cosmo_likelihood.py
+++ b/hierarc/Likelihood/cosmo_likelihood.py
@@ -91,6 +91,10 @@ def __init__(
             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(
diff --git a/test/test_Likelihood/test_cosmo_likelihood.py b/test/test_Likelihood/test_cosmo_likelihood.py
index a6abaf36..ddd85cb6 100644
--- a/test/test_Likelihood/test_cosmo_likelihood.py
+++ b/test/test_Likelihood/test_cosmo_likelihood.py
@@ -261,6 +261,17 @@ def test_bao_likelihood_integration(self):
         logl = cosmoL.likelihood(args=args, verbose=True)
         assert logl < 0
 
+        with pytest.raises(NotImplementedError):
+            cosmoL = CosmoLikelihood(
+                [],
+                self.cosmology,
+                kwargs_model,
+                self.kwargs_bounds,
+                bao_likelihood="DESI_DR2",
+                interpolate_cosmo=True,
+                cosmo_fixed=None,
+            )
+
     def test_kde_likelihood_integration(self):
         chain = import_Planck_chain(
             _PATH_2_PLANCKDATA,

From 03e0655dd967e231af6bae52b711b37e533f4882 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 22 May 2025 13:12:15 +0000
Subject: [PATCH 10/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 .../test_BAOLikelihood/test_bao_likelihood.py               | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
index 205797da..17a8062a 100644
--- a/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
+++ b/test/test_Likelihood/test_BAOLikelihood/test_bao_likelihood.py
@@ -34,7 +34,9 @@ def setup_method(self):
             "cov": cov,
         }
 
-        self.likelihood = BAOLikelihood(sample_name="CUSTOM", **self.kwargs_bao_likelihood)
+        self.likelihood = BAOLikelihood(
+            sample_name="CUSTOM", **self.kwargs_bao_likelihood
+        )
         self.dists_true = dist_true
         self.rd_true = rd
         self.sigma_d_true = sigma_d
@@ -84,7 +86,7 @@ def test_raise(self):
             self.likelihood.log_likelihood(self.cosmo_true)
 
         kwargs_bao_likelihood = self.kwargs_bao_likelihood.copy()
-        kwargs_bao_likelihood["distance_type"] = ['Unknown'] * self.num
+        kwargs_bao_likelihood["distance_type"] = ["Unknown"] * self.num
         test_L = BAOLikelihood(sample_name="CUSTOM", **kwargs_bao_likelihood)
 
         with pytest.raises(ValueError):