Quick start
===========

This example mirrors what ``scripts/measure_joint_sumstat.py`` does for a
single BGS VLIM M★ sample.

.. code-block:: python

    import numpy as np
    from pathlib import Path
    from astropy.cosmology import FlatLambdaCDM
    import sum_stat as ss
    from sum_stat.covariance.jackknife import (
        assign_jackknife_regions,
        jackknife_covariance_from_subsamples,
    )

    # ── 1. Cosmology and bin definitions ──────────────────────────────────────
    cosmo = FlatLambdaCDM(H0=67.4, Om0=0.315, Ob0=0.049, name="Planck18")

    MSTAR_BINS = np.arange(9.0, 12.6, 0.25)                          # log10(M*/Msun), 0.25 dex
    RP_BINS    = np.logspace(np.log10(0.01), np.log10(60.0), 31)     # Mpc, 30 bins
    PI_MAX     = 100.0                                                 # Mpc
    THETA_BINS = np.logspace(np.log10(0.01), np.log10(60.0), 31)     # arcmin, 30 bins
    N_JK       = 100

    # ── 2. Load catalogues ───────────────────────────────────────────────────
    # Replace with your actual loading code; e.g. from a FITS file:
    #   from astropy.io import fits
    #   d = fits.open("LS10_VLIM_ANY_10.0_..._DATA.fits")[1].data
    #   gal = ss.GalaxyCatalogue(ra=d["RA"], dec=d["DEC"], redshift=d["BEST_Z"],
    #                            log10_mstar=d["LPH_MASS_BEST"], weight=np.ones(len(d)))
    rng = np.random.default_rng(42)
    N_GAL, N_RAND = 50_000, 250_000

    gal = ss.GalaxyCatalogue(
        ra=rng.uniform(0, 60, N_GAL),
        dec=rng.uniform(-6, 6, N_GAL),
        redshift=rng.uniform(0.05, 0.18, N_GAL),
        log10_mstar=rng.normal(10.5, 0.5, N_GAL),
        weight=np.ones(N_GAL),
    )
    rand = ss.GalaxyCatalogue(
        ra=rng.uniform(0, 60, N_RAND),
        dec=rng.uniform(-6, 6, N_RAND),
        redshift=rng.uniform(0.05, 0.18, N_RAND),
    )

    # Survey area from randoms via HEALPix
    import healpy as hp
    AREA_NSIDE = 64
    pix       = hp.ang2pix(AREA_NSIDE, rand.ra, rand.dec, lonlat=True)
    area_deg2 = len(np.unique(pix)) * hp.nside2pixarea(AREA_NSIDE, degrees=True)
    area_sr   = area_deg2 * (np.pi / 180.0) ** 2
    z_min, z_max = 0.05, 0.18

    # ── 3. Angular two-point correlation function w(θ) ───────────────────────
    theta_c, w, var_w = ss.w_theta(gal, rand, THETA_BINS)

    # ── 4. Projected two-point correlation function w_p(r_p) ─────────────────
    rp_c, wp, var_wp = ss.wp(gal, rand, RP_BINS, PI_MAX, cosmo)

    # ── 5. Stellar mass function — 1/Vmax with jackknife covariance ──────────
    mstar_c, phi_vmax, phi_err_vmax = ss.stellar_mass_function_vmax(
        gal, z_min, z_max, MSTAR_BINS, area_sr, cosmo,
    )

    jk_labels = assign_jackknife_regions(gal.ra, gal.dec, n_regions=N_JK)
    jk_subsamples = []
    for k in range(N_JK):
        mask = jk_labels != k
        g_k  = ss.GalaxyCatalogue(
            ra=gal.ra[mask], dec=gal.dec[mask],
            redshift=gal.redshift[mask],
            log10_mstar=gal.log10_mstar[mask],
            weight=gal.weight[mask],
        )
        _, phi_k, _ = ss.stellar_mass_function_vmax(
            g_k, z_min, z_max, MSTAR_BINS, area_sr, cosmo)
        jk_subsamples.append(phi_k)
    _, cov_vmax = jackknife_covariance_from_subsamples(np.array(jk_subsamples))

    # ── 6. Save all results to HDF5 ──────────────────────────────────────────
    out  = Path("results.h5")
    meta = dict(
        z_min=z_min, z_max=z_max, area_deg2=area_deg2, area_sr=area_sr,
        survey="bgs", weight_variant="uniform", cov_method="jackknife",
    )

    with ss.SummaryStatWriter(str(out), mode="w") as wr:
        wr.write_twopcf(
            "wtheta", theta_c, w, np.diag(var_w), THETA_BINS,
            "landy-szalay", cosmo,
            {**meta, "sep_units": "arcmin", "cov_method": "treecorr_variance"},
            sep_unit="arcmin", xi_unit="dimensionless",
        )
        wr.write_twopcf(
            "wp", rp_c, wp, np.diag(var_wp), RP_BINS,
            "landy-szalay", cosmo,
            {**meta, "pi_max_Mpc": PI_MAX, "cov_method": "treecorr_variance"},
            sep_unit="Mpc", xi_unit="Mpc",
        )
        wr.write_smf(
            "smf", mstar_c, phi_vmax, phi_err_vmax, cov_vmax,
            MSTAR_BINS, cosmo,
            {**meta, "estimator": "vmax", "n_jackknife": N_JK},
        )

    # ── 7. Read back ─────────────────────────────────────────────────────────
    with ss.SummaryStatReader(str(out)) as rd:
        smf    = rd.read_smf("smf")
        twopcf = rd.read_twopcf("wp")
        print("phi shape:", smf["phi"].shape)
        print("wp shape:",  twopcf["xi"].shape)
        print("Cosmology:", rd.cosmology("smf"))