Merged harmonic and configuration space covariance to fix bug in inference run (Cosmosis). Added a prior file for harmonic space mocks.

Author: sachaguer

cosmo_inference/scripts/cosmosis_fitting.py

@@ -167,33 +167,6 @@ def cl_to_fits(ell, cl_ee, cl_bb):
     return cl_ee_hdu, cl_bb_hdu
 
 
-def cov_cl_to_fits(cov_file, cov_hdu="COVAR_FULL"):
-    """Convert pseudo-C_ell covariance to a CosmoSIS ImageHDU."""
-    if cov_file.endswith(".fits"):
-        with fits.open(cov_file) as hdul:
-            cov_data = np.asarray(hdul[cov_hdu].data, dtype=np.float64)
-    elif cov_file.endswith(".npy"):
-        cov_data = np.load(cov_file)
-    else:
-        raise NotImplementedError(f"Unsupported pseudo-Cl covariance format: {cov_file}")
-
-    if cov_data.shape[0] != cov_data.shape[1]:
-        raise ValueError("Pseudo-Cl covariance matrix must be square")
-
-    cov_hdu = fits.ImageHDU(cov_data, name="COVMAT_CELL")
-    cov_dict = {
-        "COVDATA": "True",
-        "EXTNAME": "COVMAT_CELL",
-        "NAME_0": "CELL_EE",
-        "STRT_0": 0,
-    }
-
-    for key, value in cov_dict.items():
-        cov_hdu.header[key] = value
-
-    return cov_hdu
-
-
 def tau_to_fits(filename, theta=None):
     """
     Convert tau statistics to FITS format.
@@ -238,12 +211,16 @@ def rho_to_fits(filename, theta=None):
     return rho_stat_hdu
 
 
-def covdat_to_fits(filename_cov_xi, filename_cov_tau=None):
+def covdat_to_fits(filename_cov_xi, filename_cov_tau=None, filename_cov_cl=None, cov_cl_hdu=None):
     """
     Convert CosmoCov covariance matrix to FITS format.
 
     If tau covariance provided, block with xi covariance.
     """
+    if filename_cov_cl is not None:
+        assert cov_cl_hdu is not None, "Covariance HDU must be provided with cov_cl file."
+
+    # Load xi covariance
     covmat_xi = np.loadtxt(filename_cov_xi)
 
     if filename_cov_tau is not None:
@@ -259,6 +236,28 @@ def covdat_to_fits(filename_cov_xi, filename_cov_tau=None):
     else:
         covmat = covmat_xi
 
+    covmat_real_space_dim = covmat.shape[0]
+    # Update the covariance with the pseudo-Cl part if provided
+    if filename_cov_cl is not None:
+        if filename_cov_cl.endswith(".fits"):
+            with fits.open(filename_cov_cl) as hdul:
+                covmat_cl = np.asarray(hdul[cov_cl_hdu].data, dtype=np.float64)
+        elif filename_cov_cl.endswith(".npy"):
+            covmat_cl = np.load(filename_cov_cl)
+        else:
+            raise NotImplementedError(f"Unsupported pseudo-Cl covariance format: {filename_cov_cl}")
+
+        if covmat_cl.shape[0] != covmat_cl.shape[1]:
+            raise ValueError("Pseudo-Cl covariance matrix must be square")
+        
+        covmat = np.block(
+            [
+                [covmat, np.zeros((len(covmat), len(covmat_cl)))],
+                [np.zeros((len(covmat_cl), len(covmat))), covmat_cl],
+            ]
+        )
+    
+
     if len(covmat) != len(covmat[0]):
         raise RuntimeError("Covariance matrix is not square")
 
@@ -279,6 +278,17 @@ def covdat_to_fits(filename_cov_xi, filename_cov_tau=None):
             "NAME_3": "TAU_2_PLUS",
             "STRT_3": len(covmat_xi) + int(len(covmat_tau) / 2),
         })
+    
+    if filename_cov_tau is None:
+        filename_cov_cl and cov_dict.update({
+            "NAME_2": "CELL_EE",
+            "STRT_2": covmat_real_space_dim,
+        })
+    else:
+        filename_cov_cl and cov_dict.update({
+            "NAME_4": "CELL_EE",
+            "STRT_4": covmat_real_space_dim,
+        })
 
     for key, value in cov_dict.items():
         cov_hdu.header[key] = value
@@ -379,7 +389,7 @@ def _generate_ini_file(
             "\ndata_sets=CELL_EE"
         )
 
-    covmat_line = "covmat_name=COVMAT_CELL" if is_harmonic else "covmat_name=COVMAT"
+    covmat_line = "covmat_name=COVMAT"
 
     modifications.append((r"^\[2pt_like\]", like_section))
     modifications.append((r"^covmat_name=.*", covmat_line))
@@ -436,11 +446,17 @@ def generate_cosmosis_config(args):
 
     if args.cl_file:
         template_base_harmonic = "cosmosis_pipeline_A_ia_cell.ini"
+
+        values_file_cell = "values_ia.ini"
+        if args.mock:
+            priors_file_cell = "priors_mock_cell.ini"
+        else:
+            priors_file_cell = "priors.ini"
         _generate_ini_file(
             args,
             template_base_harmonic,
-            priors_file,
-            values_file,
+            priors_file_cell,
+            values_file_cell,
             suffix="_cell",
             is_harmonic=True,
         )
@@ -643,17 +659,6 @@ def parse_args():
         nz_hdu = nz_to_fits(args.nz_file)
         print("Loaded n(z)")
 
-        cl_ee_hdu = None
-        cl_bb_hdu = None
-        cov_cl_hdu = None
-        if args.cl_file:
-            print("Loading Cl data...")
-            ell, cl_ee, cl_bb = load_pseudo_cl(args.cl_file)
-            print(f"Loaded Cl: {len(ell)} multipoles")
-            cl_ee_hdu, cl_bb_hdu = cl_to_fits(ell, cl_ee, cl_bb)
-            cov_cl_hdu = cov_cl_to_fits(args.cov_cl, cov_hdu="COVAR_FULL")
-            print("Loaded pseudo-Cl covariance")
-
         rho_hdu = None
         tau_0_p_hdu = None
         tau_2_p_hdu = None
@@ -673,17 +678,38 @@ def parse_args():
 
             cov_hdu = covdat_to_fits(args.cov_xi, filename_cov_tau=args.cov_tau)
 
+        cl_ee_hdu = None
+        cl_bb_hdu = None
+        if args.cl_file:
+            print("Loading Cl data...")
+            ell, cl_ee, cl_bb = load_pseudo_cl(args.cl_file)
+            print(f"Loaded Cl: {len(ell)} multipoles")
+            cl_ee_hdu, cl_bb_hdu = cl_to_fits(ell, cl_ee, cl_bb)
+            if args.use_rho_tau:
+                cov_hdu = covdat_to_fits(
+                    args.cov_xi,
+                    filename_cov_tau=args.cov_tau,
+                    filename_cov_cl=args.cov_cl,
+                    cov_cl_hdu="COVAR_FULL"
+                )
+            else:
+                cov_hdu = covdat_to_fits(
+                    args.cov_xi,
+                    filename_cov_tau=None,
+                    filename_cov_cl=args.cov_cl,
+                    cov_cl_hdu="COVAR_FULL"
+                )
+            print("Loaded pseudo-Cl covariance")
+
         pri_hdr = fits.Header()
         pri_hdu = fits.PrimaryHDU(header=pri_hdr)
 
         print("Assembling FITS file...")
         hdu_list = [pri_hdu, cov_hdu]
-        if cov_cl_hdu is not None:
-            hdu_list.append(cov_cl_hdu)
         hdu_list.extend([nz_hdu, xip_hdu, xim_hdu])
 
         if args.cl_file:
-            hdu_list.extend([cl_ee_hdu, cl_bb_hdu])
+            hdu_list.extend([cl_ee_hdu])
 
         if args.use_rho_tau:
             hdu_list.extend([tau_0_p_hdu, tau_2_p_hdu, rho_hdu])