fetch_chains.py

"""Fetch the public posterior chains used by the chain-derived probes.

These chain files are NOT redistributable through this repository for
size and licensing reasons; the user must download them from the
collaborations' public archives. This script automates the download
where the URL is stable + public, and prints clear instructions for
the cases that require institutional credentials or accept-license
clickthrough.

Run:  python -m src.fetch_chains [--probe planck|act|kids|des|desi|pr4|all]

After running, the per-probe directories under data/ will contain the
chain files needed by src.load_chains.* (see download_summary.json
for an inventory of what landed and what is still pending).

Important: the chain files vary in size from ~50 MB (DESI Y1 BAO) to
~2 GB (Planck PR4 NPIPE). Download time is dominated by the Planck
products. We do not retry; if the server times out, re-run the
script and only the missing files will be fetched.
"""

from __future__ import annotations

import argparse
import json
import sys
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"


# --- Stable public URLs -----------------------------------------------------
# Each entry: (target_relative_path, url, expected_min_bytes, notes)
PUBLIC_URLS = {
    "planck_pr3": [
        # Planck PR3 base ΛCDM TT,TE,EE+lowE chains (cosmomc format,
        # 4 walker files + .paramnames + .ranges)
        ("planck2018_chains/base/plikHM_TTTEEE_lowl_lowE/"
         "base_plikHM_TTTEEE_lowl_lowE_1.txt",
         "https://irsa.ipac.caltech.edu/data/Planck/release_3/ancillary-data/"
         "cosmoparams/COM_CosmoParams_base-plikHM-TTTEEE-lowl-lowE_R3.00.zip",
         100_000_000,
         "ZIP archive; extract under data/planck2018_chains/ "
         "(Planck Legacy Archive)."),
    ],
    "act_dr6": [
        ("act_dr6_chains/actbase_lcdm_camb/actbase_lcdm_camb/"
         "actbase_lcdm_camb.1.txt",
         "https://lambda.gsfc.nasa.gov/data/suborbital/ACT/ACT_dr6/"
         "likelihood/data/v1.0/MCMCChains/actbase_lcdm_camb.tar.gz",
         50_000_000,
         "TAR archive; extract under data/act_dr6_chains/ "
         "(LAMBDA / NASA ACT DR6 release)."),
    ],
    "kids1000": [
        ("kids1000/KiDS1000_cosmis_shear_data_release/"
         "chains_and_config_files/main_chains_iterative_covariance/"
         "cosebis/chain/output_multinest_C.txt",
         "https://kids.strw.leidenuniv.nl/cs2018/KiDS1000_cosmis_shear_data_release.tar.gz",
         200_000_000,
         "TAR archive; extract under data/kids1000/ (KiDS-1000 release)."),
    ],
    "des_y3": [
        ("des_y3/chain_3x2pt_lcdm_SR_maglim.txt",
         "https://desdr-server.ncsa.illinois.edu/despublic/y3a2_files/"
         "chains/chain_3x2pt_lcdm_SR_maglim.txt",
         5_000_000,
         "Direct .txt download; ~30 MB (DES Y3 release)."),
    ],
    "desi_y1": [
        # DESI Y1 BAO + BBN baseline chain (placeholder; fill in real URL
        # from DESI DR1 public archive as available).
        ("desi_y1_chains/desi_y1_bao_bbn.txt",
         "https://data.desi.lbl.gov/public/dr1/cosmology/chains/"
         "desi_y1_bao_bbn.txt",
         1_000_000,
         "Direct .txt; small (~5 MB)."),
    ],
    "planck_pr4": [
        ("planck_pr4_chains/hlpTTTEEE_lowT_lowE_1.txt",
         "https://portal.nersc.gov/cfs/cmb/planck2020/chains/"
         "PR4_NPIPE_TTTEEE_lowT_lowE.tar.gz",
         500_000_000,
         "TAR; extract under data/planck_pr4_chains/ "
         "(NERSC; PR4 NPIPE / Tristram et al. 2024)."),
    ],
}


def _download(url: str, target: Path) -> bool:
    """Stream a file; return True on success. Skip if already present."""
    if target.exists() and target.stat().st_size > 0:
        return True
    target.parent.mkdir(parents=True, exist_ok=True)
    try:
        print(f"  fetching {url}")
        with urllib.request.urlopen(url, timeout=60) as r, open(target, "wb") as f:
            while True:
                chunk = r.read(1 << 20)
                if not chunk:
                    break
                f.write(chunk)
        return True
    except Exception as e:
        print(f"  FAIL: {e}")
        if target.exists() and target.stat().st_size == 0:
            target.unlink()
        return False


def fetch_one(probe: str) -> dict:
    if probe not in PUBLIC_URLS:
        raise KeyError(f"unknown probe '{probe}'")
    out = {"probe": probe, "fetched": [], "missing": [], "notes": []}
    for rel, url, _, note in PUBLIC_URLS[probe]:
        target = DATA / rel
        if target.exists() and target.stat().st_size > 0:
            out["fetched"].append(rel)
            continue
        ok = _download(url, target)
        (out["fetched"] if ok else out["missing"]).append(rel)
        if note and not ok:
            out["notes"].append(note)
    return out


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--probe", default="all",
                   choices=list(PUBLIC_URLS) + ["all"])
    args = p.parse_args()
    targets = list(PUBLIC_URLS) if args.probe == "all" else [args.probe]
    summary = []
    for t in targets:
        print(f"=== {t} ===")
        summary.append(fetch_one(t))
    out = DATA / "download_summary.json"
    out.write_text(json.dumps(summary, indent=2))
    n_ok = sum(len(s["fetched"]) for s in summary)
    n_missing = sum(len(s["missing"]) for s in summary)
    print(f"\nFetched {n_ok}; missing {n_missing}.")
    print(f"Wrote {out}")
    if n_missing:
        print("\nManual download required for missing files. Notes:")
        for s in summary:
            for note in s["notes"]:
                print(f"  [{s['probe']}] {note}")
        sys.exit(1)


if __name__ == "__main__":
    main()