generate_pantheon_module_from_csv.py

"""
Read a cleaned CSV (z,mu,sigma_mu) and write an importable Python module:
  src/likelihoods/data/pantheon_plus.py

The module will export:
  PANTHEON_PLUS_FULL = {"z": [...], "mu": [...], "sigma_mu": [...]}
"""

from pathlib import Path
import pandas as pd
import sys

CSV = Path("data/processed/pantheon_plus.csv")
OUT = Path("src/likelihoods/data/pantheon_plus.py")

if not CSV.exists():
    print(f"Error: {CSV} not found. Run the converter first.", file=sys.stderr)
    sys.exit(1)

# Read CSV robustly; allow for a truncated final line
try:
    df = pd.read_csv(CSV, usecols=["z", "mu", "sigma_mu"])
except ValueError:
    # If header or column names differ, try to read all and pick first three columns
    df = pd.read_csv(CSV)
    if set(["z", "mu", "sigma_mu"]).issubset(df.columns):
        df = df[["z", "mu", "sigma_mu"]]
    else:
        # fallback: assume first three columns are z, mu, sigma_mu
        df = df.iloc[:, :3]
        df.columns = ["z", "mu", "sigma_mu"]

# Coerce to numeric and drop invalid rows
for c in ["z", "mu", "sigma_mu"]:
    df[c] = pd.to_numeric(df[c], errors="coerce")
df = df.dropna().reset_index(drop=True)

if df.shape[0] == 0:
    print("Error: no valid numeric rows found in CSV.", file=sys.stderr)
    sys.exit(2)

# Convert to plain Python lists of floats
z = df["z"].astype(float).tolist()
mu = df["mu"].astype(float).tolist()
sigma = df["sigma_mu"].astype(float).tolist()

# Write canonical module
OUT.parent.mkdir(parents=True, exist_ok=True)
with open(OUT, "w") as f:
    f.write("# Auto-generated Pantheon+ dataset (numeric lists)\n")
    f.write("# Generated by scripts/generate_pantheon_module_from_csv.py\n\n")
    f.write("PANTHEON_PLUS_FULL = {\n")
    f.write(f"    'z': {z},\n")
    f.write(f"    'mu': {mu},\n")
    f.write(f"    'sigma_mu': {sigma}\n")
    f.write("}\n")

print(f"Wrote {OUT} with {len(z)} SNe")