S.T.A.R.-Program/scripts/build_cosmology_data.py at main · LcosmosS/S.T.A.R.-Program · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
#!/usr/bin/env python3
import zipfile
import numpy as np
import pandas as pd
from pathlib import Path
import json

ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = ROOT / "data" / "cosmic"
OUT_DIR = ROOT / "src" / "likelihoods" / "data"
OUT_DIR.mkdir(parents=True, exist_ok=True)

PLANCK_CHAIN_ZIP = DATA_DIR / "COM_CosmoParams_base-plikHM-TTTEEE-lowl-lowE_R3.00.zip"
BAO_ZIP = DATA_DIR / "bao_data-master.zip"


def write_module(name, obj):
    path = OUT_DIR / f"{name}.py"
    with open(path, "w") as f:
        f.write("# Auto-generated by CI. Do not edit.\n")
        f.write(f"{name.upper()} = ")
        json.dump(obj, f, indent=2)
        f.write("\n")


# -----------------------------
# 1. Extract BAO + CC data
# -----------------------------
with zipfile.ZipFile(BAO_ZIP, "r") as z:
    z.extractall(DATA_DIR / "bao_data-master")

cc_path = DATA_DIR / "cc" / "Hz_Moresco2022.dat"
bao_dir = DATA_DIR / "bao_data-master" / "BAO"

# Cosmic Chronometers
cc = np.loadtxt(cc_path)
COSMIC_CHRONOMETERS = {
    "z": cc[:, 0].tolist(),
    "H": cc[:, 1].tolist(),
    "sigma_H": cc[:, 2].tolist(),
}
write_module("cosmic_chronometers", COSMIC_CHRONOMETERS)

# DESI BAO DR1
bao_files = sorted(bao_dir.glob("*.dat"))
z_list = []
DM_list = []
DM_err = []
Hrd_list = []
Hrd_err = []

for bf in bao_files:
    arr = np.loadtxt(bf)
    z_list.append(arr[0])
    DM_list.append(arr[1])
    DM_err.append(arr[2])
    Hrd_list.append(arr[3])
    Hrd_err.append(arr[4])

DESI_BAO_DR1 = {
    "z": z_list,
    "DM_over_rd": DM_list,
    "sigma_DM": DM_err,
    "H_rd": Hrd_list,
    "sigma_H": Hrd_err,
}
write_module("desi_bao_dr1", DESI_BAO_DR1)

# -----------------------------
# 2. Planck 2015 distance priors
# -----------------------------
PLANCK_2015_PRIORS = {
    "R": 1.74963,
    "lA": 301.80845,
    "ombh2": 0.02226,
    "cov": [
        [0.00053, 0.00211, -0.000001],
        [0.00211, 0.08912, -0.000012],
        [-0.000001, -0.000012, 0.00000004],
    ],
}
write_module("planck_2015", PLANCK_2015_PRIORS)

# -----------------------------
# 3. Reconstruct Planck 2018 priors
# -----------------------------
with zipfile.ZipFile(PLANCK_CHAIN_ZIP, "r") as z:
    z.extractall(DATA_DIR / "planck_chain")

chain_dir = DATA_DIR / "planck_chain"
samples = []

# Load only the columns we need: omegabh2, omegach2, H0, theta_MC
usecols = (2, 3, 4, 5)

for f in chain_dir.rglob("*.txt"):
    try:
        arr = np.loadtxt(f, usecols=usecols)
        samples.append(arr)
    except Exception:
        pass  # skip non-chain files

if len(samples) == 0:
    raise RuntimeError("No valid Planck chain files found.")

samples = np.vstack(samples)

ombh2 = samples[:, 0]
omegach2 = samples[:, 1]
H0 = samples[:, 2]
theta = samples[:, 3]

# Compute R and lA (approximate)
R = np.sqrt((ombh2 + omegach2) * (H0 / 100) ** 2) * theta
lA = np.pi / theta

PLANCK_2018_RECON_PRIORS = {
    "R": float(np.mean(R)),
    "lA": float(np.mean(lA)),
    "ombh2": float(np.mean(ombh2)),
    "cov": np.cov(np.vstack([R, lA, ombh2])).tolist(),
}

write_module("planck_2018_recon", PLANCK_2018_RECON_PRIORS)

print("Cosmology data modules generated.")