py-lfkit 0.1.4__py3-none-any.whl

lfkit/corrections/kcorrect_grids.py

@@ -0,0 +1,242 @@
+"""kcorrect k(z) grid generation and interpolation utilities.
+
+This module builds tables of k(z) values for one or more **anchors** on a
+specified redshift grid and provides helpers to interpolate those tables.
+
+An *anchor* is simply a label associated with a vector of kcorrect template
+coefficients. Each coefficient vector represents a particular template
+mixture and therefore defines a specific spectral energy distribution.
+Evaluating kcorrect with those coefficients produces the corresponding
+k(z) curve for a given set of filter responses.
+
+The main purpose of this module is to turn a small set of anchors into
+reusable k(z) grids that can be cached and quickly interpolated during
+analysis. The resulting tables store k(z) as a function of redshift for
+each anchor and response band, allowing fast lookup without repeatedly
+calling the kcorrect solver.
+
+Anchors intentionally carry **no semantic meaning** here. They are not
+interpreted as galaxy types or populations; they are simply convenient
+labels for different template mixtures used to generate k(z) curves.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from lfkit.utils.interpolation import Interpolator, build_1d_interpolator
+
+from .kcorrect_backend import build_kcorrect
+
+
+def compute_k_table(
+    *,
+    kc,
+    z_grid: np.ndarray,
+    coeffs_by_anchor: dict[str, np.ndarray],
+    band_shift: float | None = None,
+    anchor_z0: bool = True,
+) -> dict[str, np.ndarray]:
+    """Compute k(z) values on a redshift grid for each anchor.
+
+    This function evaluates k-corrections on a fixed redshift grid using a set
+    of precomputed template coefficient vectors (“anchors”). Each anchor
+    represents a particular SED mixture and therefore defines a specific
+    k(z) curve. The output is a table of k(z) values for every anchor across
+    the provided grid, which can later be used for interpolation or lookup.
+
+    By default the table is normalized so that k(z=0)=0 for each band. This
+    convention is commonly used when working with k-corrections because it
+    removes arbitrary offsets and ensures that all curves are anchored to
+    the same reference point.
+    """
+    z = np.asarray(z_grid, float)
+    if z.ndim != 1 or z.size < 2 or np.any(~np.isfinite(z)):
+        raise ValueError("z_grid must be a finite 1D array with >=2 points.")
+
+    out: dict[str, np.ndarray] = {}
+
+    for label, coeffs in coeffs_by_anchor.items():
+        c = np.asarray(coeffs, float).reshape(-1)
+
+        rows: list[np.ndarray] = []
+        nband: int | None = None
+
+        for zi in z:
+            zi = float(zi)
+
+            if zi == 0.0 and anchor_z0:
+                if nband is None:
+                    if band_shift is None:
+                        test = kc.kcorrect(redshift=1e-6, coeffs=c)
+                    else:
+                        test = kc.kcorrect(
+                            redshift=1e-6, coeffs=c, band_shift=float(band_shift)
+                        )
+                    nband = int(np.asarray(test, float).size)
+                kcorr = np.zeros(nband, dtype=float)
+            else:
+                if band_shift is None:
+                    kcorr = kc.kcorrect(redshift=zi, coeffs=c)
+                else:
+                    kcorr = kc.kcorrect(
+                        redshift=zi, coeffs=c, band_shift=float(band_shift)
+                    )
+                kcorr = np.asarray(kcorr, float)
+                if nband is None:
+                    nband = int(kcorr.size)
+
+            rows.append(np.asarray(kcorr, float))
+
+        karr = np.vstack(rows)
+        if karr.shape[0] != z.size:
+            raise ValueError(
+                f"BUG: built K with Nz={karr.shape[0]} but z has Nz={z.size}"
+            )
+
+        if anchor_z0:
+            if not np.any(np.isfinite(karr)):
+                raise ValueError(f"All-NaN K grid for anchor={label!r}.")
+            karr = karr - karr[0:1, :]
+
+        out[str(label)] = karr
+
+    return out
+
+
+def build_kcorr_grid_package(
+    *,
+    responses_in: list[str],
+    responses_out: list[str] | None,
+    responses_map: list[str] | None,
+    coeffs_by_anchor: dict[str, np.ndarray],
+    z_grid: np.ndarray,
+    band_shift: float | None = 0.1,
+    response_dir: str | None = None,
+    redshift_range: tuple[float, float] = (0.0, 3.5),
+    nredshift: int = 4000,
+) -> dict[str, Any]:
+    """Build a packaged k(z) grid for a set of anchors.
+
+    This function generates a self-contained data package holding k(z)
+    tables evaluated on a common redshift grid for multiple anchors.
+    Each anchor corresponds to a template mixture that defines a specific
+    spectral energy distribution and therefore a specific k(z) curve.
+
+    The resulting package contains the redshift grid, the computed k(z)
+    tables, the filter responses used, and basic metadata describing the
+    setup. It is intended to serve as a reusable object that can be saved,
+    cached, or passed to interpolation routines without recomputing the
+    k-corrections.
+    """
+    if responses_out is None:
+        responses_out = responses_in
+    if responses_map is None:
+        responses_map = responses_in
+
+    kc = build_kcorrect(
+        responses_in=responses_in,
+        responses_out=responses_out,
+        responses_map=responses_map,
+        response_dir=response_dir,
+        redshift_range=redshift_range,
+        nredshift=nredshift,
+    )
+
+    nt = int(kc.templates.restframe_flux.shape[0])
+
+    for label, c in coeffs_by_anchor.items():
+        c = np.asarray(c, float)
+        if c.shape != (nt,):
+            raise ValueError(f"{label}: coeff shape {c.shape} != ({nt},)")
+        if not np.all(np.isfinite(c)):
+            raise ValueError(f"{label}: coeffs contain non-finite values.")
+        if np.any(c < 0):
+            raise ValueError(f"{label}: negative coeffs not allowed.")
+        if float(np.sum(c)) <= 0:
+            raise ValueError(f"{label}: coeffs sum <= 0.")
+
+    z_grid = np.asarray(z_grid, float)
+    K = compute_k_table(
+        kc=kc,
+        z_grid=z_grid,
+        coeffs_by_anchor=coeffs_by_anchor,
+        band_shift=band_shift,
+        anchor_z0=True,
+    )
+
+    return dict(
+        meta=dict(
+            backend="kcorrect",
+            band_shift=band_shift,
+            redshift_range=(float(z_grid[0]), float(z_grid[-1])),
+            nredshift=int(nredshift),
+            response_dir=str(response_dir) if response_dir is not None else None,
+        ),
+        z=z_grid,
+        responses_in=list(map(str, responses_in)),
+        responses_out=list(map(str, responses_out)),
+        responses_map=list(map(str, responses_map)),
+        anchors=sorted(list(map(str, coeffs_by_anchor.keys()))),
+        K=K,  # mapping: anchor_label -> (Nz, Nband)
+    )
+
+
+def kcorr_interpolators(
+    pkg: dict[str, Any],
+    *,
+    method: str = "pchip",
+    extrapolate: bool = True,
+) -> dict[str, dict[str, Interpolator | None]]:
+    """Create interpolation functions for k(z).
+
+    This function converts a precomputed k(z) grid package into a set of
+    interpolators that return k(z) for any redshift within the grid range.
+    An interpolator is produced for each combination of anchor and output
+    response band.
+
+    The resulting structure allows k-corrections to be evaluated quickly
+    without recomputing the underlying tables, making it convenient to
+    integrate precomputed grids into analysis workflows that repeatedly
+    query k(z) values.
+    """
+    z = np.asarray(pkg["z"], float)
+    responses_out = list(pkg["responses_out"])
+    anchors = list(pkg["anchors"])
+
+    out: dict[str, dict[str, Interpolator | None]] = {}
+    for label in anchors:
+        ktz = np.asarray(pkg["K"][label], float)  # (Nz, Nband)
+
+        if ktz.shape[0] != z.size:
+            raise ValueError(
+                f"Shape mismatch for anchor={label!r}:"
+                f"K has Nz={ktz.shape[0]} vs z={z.size}."
+            )
+        if ktz.shape[1] != len(responses_out):
+            raise ValueError(
+                f"Shape mismatch for anchor={label!r}: "
+                f"K has Nband={ktz.shape[1]} "
+                f"vs responses_out={len(responses_out)}."
+            )
+
+        out[str(label)] = {}
+        for j, band in enumerate(responses_out):
+            y = np.asarray(ktz[:, j], float)
+            ok = np.isfinite(z) & np.isfinite(y)
+
+            if np.count_nonzero(ok) < 2:
+                out[str(label)][str(band)] = None
+                continue
+
+            out[str(label)][str(band)] = build_1d_interpolator(
+                z[ok],
+                y[ok],
+                method=method,
+                extrapolate=extrapolate,
+                extrap_mode="linear_tail",
+            )
+
+    return out

lfkit/corrections/poggianti1997.py

@@ -0,0 +1,386 @@
+"""Poggianti (1997) k- and e-correction tables and interpolators.
+
+Loads Poggianti (1997) k- and e-correction curves from CSV tables, builds
+interpolators for those curves (PCHIP, Akima, or linear), and optionally remaps
+the e-correction redshift grid to a target cosmology by matching lookback time.
+
+Table I/O and parsing live in ``lfkit.utils.io``.
+Interpolation utilities live in ``lfkit.utils.interpolation``.
+Unit conversions live in ``lfkit.utils.units``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+from lfkit.cosmo.cosmology import cosmo_object as build_cosmo, lookback_time_gyr
+from lfkit.utils.interpolation import (
+    InterpMethod,
+    Interpolator,
+    build_1d_interpolator,
+    prep_strictly_increasing_xy,
+)
+from lfkit.utils.io import (
+    POGGIANTI1997_PKG,
+    available_from_table,
+    extract_series,
+    load_vizier_csv,
+    resolve_packaged_csv,
+)
+from lfkit.utils.units import h0_km_s_mpc_to_gyr_inv
+
+__all__ = (
+    "available_pairs",
+    "load_poggianti1997_tables",
+    "describe_poggianti1997_available",
+    "poggianti1997_time_since_bb_gyr",
+    "poggianti1997_lookback_time_gyr",
+    "z_from_lookback_time",
+    "poggianti1997_to_accelerating_redshift",
+    "make_kcorr_interpolator",
+    "make_ecorr_interpolator",
+    "extract_sed_spectrum",
+)
+
+
+def available_pairs(tab: np.ndarray, *, min_points: int = 5) -> dict[str, list[str]]:
+    """List usable (band -> SEDs) pairs in a Poggianti-style table.
+
+    Args:
+        tab: Structured array for a single correction table (k-corr or e-corr).
+        min_points: Minimum number of samples required per extracted series.
+
+    Returns:
+        Mapping from band label to a list of SED labels that have usable data.
+    """
+    bands, seds = available_from_table(tab)
+    out: dict[str, list[str]] = {b: [] for b in bands}
+    for b in bands:
+        for s in seds:
+            try:
+                extract_series(tab, band=b, sed=s, min_points=min_points)
+            except ValueError:
+                continue
+            out[b].append(s)
+    return out
+
+
+def load_poggianti1997_tables(
+    *,
+    band: str = "r",
+    sed: str = "E",
+    kcorr_path: str | Path | None = None,
+    ecorr_path: str | Path | None = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Load Poggianti (1997) k- and e-correction curves for a band and SED.
+
+    Args:
+        band: Band/filter label in the tables (e.g. ``"r"``, ``"B"``, ``"V"``).
+        sed: SED column label (e.g. ``"E"``, ``"Sa"``, ``"Sc"``).
+        kcorr_path: Optional path to ``kcorr.csv``. If not provided, the
+            packaged file is used.
+        ecorr_path: Optional path to ``ecorr.csv``. If not provided, the
+            packaged file is used.
+
+    Returns:
+        Tuple ``(z_k, kcorr, z_e, ecorr)``:
+        - ``z_k`` and ``kcorr`` are the k-correction curve.
+        - ``z_e`` and ``ecorr`` are the e-correction curve.
+
+    Raises:
+        ValueError: If the requested (band, sed) is not available or contains
+            insufficient usable samples.
+    """
+    if kcorr_path is None:
+        kcorr_path = resolve_packaged_csv("kcorr.csv", pkg=POGGIANTI1997_PKG)
+    if ecorr_path is None:
+        ecorr_path = resolve_packaged_csv("ecorr.csv", pkg=POGGIANTI1997_PKG)
+
+    ktab = load_vizier_csv(kcorr_path)
+    etab = load_vizier_csv(ecorr_path)
+
+    z_k, kcorr = extract_series(ktab, band=band, sed=sed)
+    z_e, ecorr = extract_series(etab, band=band, sed=sed)
+    return z_k, kcorr, z_e, ecorr
+
+
+def describe_poggianti1997_available(
+    *,
+    kcorr_path: str | Path | None = None,
+    ecorr_path: str | Path | None = None,
+) -> dict[str, dict[str, list[str]]]:
+    """Summarize available bands and SED columns in Poggianti CSV tables.
+
+    Args:
+        kcorr_path: Optional path to ``kcorr.csv``. If not provided, the
+            packaged file is used.
+        ecorr_path: Optional path to ``ecorr.csv``. If not provided, the
+            packaged file is used.
+
+    Returns:
+        A dictionary with keys ``"kcorr"`` and ``"ecorr"``. Each contains:
+        - ``"bands"``: list of available band labels
+        - ``"seds"``: list of available SED column labels
+    """
+    if kcorr_path is None:
+        kcorr_path = resolve_packaged_csv("kcorr.csv", pkg=POGGIANTI1997_PKG)
+    if ecorr_path is None:
+        ecorr_path = resolve_packaged_csv("ecorr.csv", pkg=POGGIANTI1997_PKG)
+
+    ktab = load_vizier_csv(kcorr_path)
+    etab = load_vizier_csv(ecorr_path)
+
+    k_bands, k_seds = available_from_table(ktab)
+    e_bands, e_seds = available_from_table(etab)
+    return {
+        "kcorr": {"bands": k_bands, "seds": k_seds},
+        "ecorr": {"bands": e_bands, "seds": e_seds},
+    }
+
+
+def poggianti1997_time_since_bb_gyr(z: np.ndarray | float) -> np.ndarray:
+    """Return cosmic time since the Big Bang in the Poggianti (1997) cosmology.
+
+    Args:
+        z: Redshift value(s).
+
+    Returns:
+        Cosmic time since the Big Bang at each redshift, in Gyr.
+
+    Notes:
+        This uses the decelerating cosmology assumed by Poggianti (1997)
+        (q0 = 0.225, H0 = 50 km/s/Mpc). It is intended for lookback-time
+        matching when remapping e-about to a different cosmology.
+    """
+    q0 = 0.225
+    h0_km_s_mpc = 50.0
+
+    z = np.asarray(z, dtype=float)
+    h0_gyr_inv = h0_km_s_mpc_to_gyr_inv(h0_km_s_mpc)
+
+    term1 = -4.0 * q0 / (h0_gyr_inv * np.power(1.0 - 2.0 * q0, 1.5))
+    root_val = np.sqrt((1.0 + 2.0 * q0 * z) / (1.0 - 2.0 * q0))
+    term2 = root_val
+    term3 = 2.0 * (1.0 - (1.0 + 2.0 * q0 * z) / (1.0 - 2.0 * q0))
+    term4 = 0.25 * np.log(np.abs((1.0 + root_val) / (1.0 - root_val)))
+
+    return term1 * (term2 / term3 + term4)
+
+
+def poggianti1997_lookback_time_gyr(z: np.ndarray | float) -> np.ndarray:
+    """Return lookback time in the Poggianti (1997) cosmology.
+
+    Args:
+        z: Redshift value(s).
+
+    Returns:
+        Lookback time to each redshift (relative to z=0), in Gyr.
+    """
+    z = np.asarray(z, dtype=float)
+    return poggianti1997_time_since_bb_gyr(0.0) - poggianti1997_time_since_bb_gyr(z)
+
+
+def _build_tlb_to_z_grid(
+    cosmo_obj, *, zmax: float, nz: int
+) -> tuple[np.ndarray, np.ndarray]:
+    """Build a monotonic mapping from lookback time to redshift for a cosmology.
+
+    Args:
+        cosmo_obj: Cosmology object accepted by ``lookback_time_gyr``.
+        zmax: Maximum redshift for the mapping grid.
+        nz: Number of samples used to build the mapping grid.
+
+    Returns:
+        Tuple ``(t_grid, z_grid)`` with strictly increasing ``t_grid``.
+    """
+    if zmax <= 0:
+        raise ValueError("zmax must be > 0.")
+    if nz < 32:
+        raise ValueError("nz is too small; use at least ~256 in practice.")
+
+    z_grid = np.linspace(0.0, float(zmax), int(nz))
+    t_grid = lookback_time_gyr(cosmo_obj, z_grid)
+
+    order = np.argsort(t_grid)
+    t_grid = t_grid[order]
+    z_grid = z_grid[order]
+
+    keep = np.ones_like(t_grid, dtype=bool)
+    keep[1:] = t_grid[1:] > t_grid[:-1]
+    return t_grid[keep], z_grid[keep]
+
+
+def z_from_lookback_time(
+    cosmo_obj,
+    t_lb_gyr: np.ndarray | float,
+    *,
+    zmax: float = 20.0,
+    nz: int = 4096,
+) -> np.ndarray:
+    """Invert lookback time to redshift using a precomputed interpolation grid.
+
+    Args:
+        cosmo_obj: Cosmology object accepted by ``lookback_time_gyr``.
+        t_lb_gyr: Lookback time value(s) in Gyr.
+        zmax: Maximum redshift used to build the inversion grid.
+        nz: Number of samples used to build the inversion grid.
+
+    Returns:
+        Redshift value(s) corresponding to ``t_lb_gyr``.
+
+    Raises:
+        ValueError: If requested lookback times fall outside the grid range.
+    """
+    t_lb = np.asarray(t_lb_gyr, dtype=float)
+    t_grid, z_grid = _build_tlb_to_z_grid(cosmo_obj, zmax=zmax, nz=nz)
+
+    if np.any(t_lb < t_grid[0]) or np.any(t_lb > t_grid[-1]):
+        raise ValueError(
+            "Target lookback time outside inversion range. "
+            "Increase zmax or check inputs."
+        )
+
+    return np.interp(t_lb, t_grid, z_grid)
+
+
+def poggianti1997_to_accelerating_redshift(
+    z_dec: np.ndarray | float,
+    *,
+    cosmo_obj,
+    zmax: float = 20.0,
+    nz: int = 4096,
+) -> np.ndarray:
+    """Map Poggianti (1997) redshifts to a target cosmology via lookback time.
+
+    Args:
+        z_dec: Redshift value(s) in the Poggianti (1997) cosmology.
+        cosmo_obj: Target cosmology used to compute lookback times.
+        zmax: Maximum redshift used to build the inversion grid.
+        nz: Number of samples used to build the inversion grid.
+
+    Returns:
+        Redshift value(s) in the target cosmology with matched lookback time.
+    """
+    z_dec = np.asarray(z_dec, dtype=float)
+    t_lb = poggianti1997_lookback_time_gyr(z_dec)
+    return z_from_lookback_time(cosmo_obj, t_lb, zmax=zmax, nz=nz)
+
+
+def make_kcorr_interpolator(
+    z_k: np.ndarray,
+    kcorr: np.ndarray,
+    *,
+    method: InterpMethod = "pchip",
+    extrapolate: bool = True,
+    z_end: float = 20.0,
+    tail: bool = True,
+) -> Interpolator:
+    """Create an interpolator for a Poggianti k-correction curve.
+
+    Args:
+        z_k: Redshift samples for the k-correction table.
+        kcorr: K-correction values at ``z_k``.
+        method: Interpolation method (``"pchip"``, ``"akima"``, or ``"linear"``).
+        extrapolate: Whether to allow evaluation outside the tabulated range.
+        z_end: Redshift endpoint for the optional high-z tail.
+        tail: Whether to append a linear high-z tail out to ``z_end``.
+
+    Returns:
+        An interpolator callable ``K(z)``.
+    """
+    z = np.r_[0.0, np.asarray(z_k, float)]
+    y = np.r_[0.0, np.asarray(kcorr, float)]
+    z, y = prep_strictly_increasing_xy(z, y)
+
+    if tail and z[-1] < z_end:
+        i0 = max(0, z.size - 3)
+        dz = z[-1] - z[i0]
+        if dz <= 0:
+            raise ValueError("Non-increasing z near the tail; cannot form slope.")
+        slope = (y[-1] - y[i0]) / dz
+        y_end = y[-1] + slope * (z_end - z[-1])
+        z = np.r_[z, z_end]
+        y = np.r_[y, y_end]
+
+    return build_1d_interpolator(z, y, method=method, extrapolate=extrapolate)
+
+
+def make_ecorr_interpolator(
+    z_e: np.ndarray,
+    ecorr: np.ndarray,
+    *,
+    original_z: bool,
+    cosmo=None,
+    zmap_zmax: float = 20.0,
+    zmap_nz: int = 4096,
+    method: InterpMethod = "pchip",
+    extrapolate: bool = True,
+) -> Interpolator:
+    """Create an interpolator for a Poggianti e-correction curve.
+
+    Args:
+        z_e: Redshift samples for the e-correction table.
+        ecorr: E-correction values at ``z_e``.
+        original_z: If True, interpret ``z_e`` as Poggianti (1997) redshifts.
+            If False, remap ``z_e`` to the target cosmology via lookback time.
+        cosmo: Target cosmology used for the remapping when ``original_z=False``.
+            If not provided, lfkit's default cosmology is used.
+        zmap_zmax: Maximum redshift used to build the remapping inversion grid.
+        zmap_nz: Number of samples used to build the remapping inversion grid.
+        method: Interpolation method (``"pchip"``, ``"akima"``, or ``"linear"``).
+        extrapolate: Whether to extrapolate beyond the tabulated domain.
+
+    Returns:
+        An interpolator callable ``E(z)``.
+    """
+    z_e = np.asarray(z_e, float)
+    ecorr = np.asarray(ecorr, float)
+
+    if not original_z:
+        if cosmo is None:
+            cosmo = build_cosmo()
+        z_e = poggianti1997_to_accelerating_redshift(
+            z_e, cosmo_obj=cosmo, zmax=zmap_zmax, nz=zmap_nz
+        )
+
+    z = np.r_[0.0, z_e]
+    y = np.r_[0.0, ecorr]
+    return build_1d_interpolator(z, y, method=method, extrapolate=extrapolate)
+
+
+def extract_sed_spectrum(
+    sed_tab: np.ndarray,
+    sed_col: str,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Extract a Poggianti sed.csv column as (wave_A, flux) in rest frame.
+
+    Args:
+        sed_tab: Table returned by lfkit.utils.io.load_vizier_csv for sed.csv.
+        sed_col: Column name in sed.csv, e.g. "logF03".
+
+    Returns:
+        wave_A: Wavelength grid in Å (sorted ascending).
+        flux: Linear flux values (10**logF) on the same grid.
+
+    Raises:
+        ValueError: If required columns are missing or sed_col not present.
+    """
+    cols = list(sed_tab.dtype.names or [])
+    if "Lam" not in cols:
+        raise ValueError(f"sed.csv missing Lam column. cols={cols}")
+    if sed_col not in cols:
+        raise ValueError(
+            f"sed.csv: sed_col={sed_col!r} not present. Available: {cols[:20]} ..."
+        )
+
+    wave_nm = np.asarray(sed_tab["Lam"], float)
+    logf = np.asarray(sed_tab[sed_col], float)
+
+    ok = np.isfinite(wave_nm) & np.isfinite(logf)
+    wave_A = 10.0 * wave_nm[ok]  # nm -> Å
+    flux = 10.0 ** logf[ok]  # log10 -> linear
+
+    order = np.argsort(wave_A)
+    return wave_A[order], flux[order]