classy-szlite 0.2.0__py3-none-any.whl

classy_szlite/__init__.py

@@ -0,0 +1,50 @@
+"""classy_szlite — fast, differentiable cosmology in pure JAX.
+
+A standalone Python package providing JIT-compiled, ``jax.grad``-friendly
+access to:
+
+  * CMB angular power spectra (TT, TE, EE)
+  * Linear and non-linear matter Pk
+  * Distances (H(z), chi, Da)
+  * Derived parameters (σ8, Ω_m, S8)
+  * Halo-model tSZ Cl^yy (Arnaud 2010 GNFW)
+
+Backed by the high-accuracy ede-v2 CosmoPower emulators (Bolliet et al.).
+LCDM-equivalent cosmology is obtained by leaving ``fEDE = 0.001`` (the
+default) — no separate emulator path needed.
+
+Runtime deps: ``jax``, ``numpy``, ``mcfit``.
+
+Quick start:
+
+>>> import classy_szlite as csl
+>>> cosmo = csl.CosmoParams()
+>>> d = csl.derived(cosmo)
+>>> ell_cls = csl.cl_TTTEEE(cosmo)
+>>> k, pk = csl.Pk(cosmo, [0., 0.5, 1., 2.])
+>>> Hz, chi, Da = csl.distances(cosmo, [0.1, 0.5, 1.0])
+>>> profile = csl.ProfileParamsA10(P0=8.13, beta=5.48, B=1.25)
+>>> import jax.numpy as jnp
+>>> cl_1h, cl_2h = csl.cl_yy(cosmo, profile, ell=jnp.geomspace(2, 9000, 80))
+>>> ev = csl.cl_yy_factory(cosmo, ell=jnp.geomspace(2, 9000, 80))  # fast path
+>>> cl_1h, cl_2h = ev(profile)
+"""
+__version__ = "0.2.0"
+
+from .params import CosmoParams, ProfileParamsA10
+from .api import (
+    derived,
+    cl_TTTEEE,
+    Pk, Pnl,
+    distances,
+    cl_yy,
+    cl_yy_factory,
+    cosmo_to_dict,
+)
+
+__all__ = [
+    "__version__",
+    "CosmoParams", "ProfileParamsA10",
+    "derived", "cl_TTTEEE", "Pk", "Pnl", "distances", "cl_yy", "cl_yy_factory",
+    "cosmo_to_dict",
+]

classy_szlite/_emulator.py

@@ -0,0 +1,180 @@
+"""Pure-JAX CosmoPower emulator forward pass.
+
+Loads ``.npz`` files produced by the CosmoPower training pipeline and
+exposes a JAX-traceable ``predict`` function. Supports two architectures:
+
+1. **Plain NN** (TT, DER, HZ, DAZ, PKL, PKNL, S8Z, etc.):
+   output = feature_train_std * NN(p_in) + feature_train_mean
+
+2. **NN + PCA** (TE and similar that decompress PCA coefficients):
+   output = ((NN(p_in) @ pca_transform_matrix) + pca_mean) * pca_std
+            * feature_train_std + feature_train_mean
+   (exact ordering inferred from cosmopower/classy_szfast convention)
+
+The NN itself uses the "α-β sigmoid" activation introduced in
+Spurio Mancini, Piras, Alsing et al. (2022, MNRAS 511, 1771) — a
+trainable smooth interpolation between identity and sigmoid:
+
+    h_out = (β + sigmoid(α * z) * (1 - β)) * z
+
+No keras / tensorflow dependency: the .npz files are read with numpy,
+and the forward pass is a handful of jnp ops.
+"""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+jax.config.update("jax_enable_x64", True)
+
+
+def _sigmoid(x):
+    return 1.0 / (1.0 + jnp.exp(-x))
+
+
+@dataclass(frozen=True)
+class Emulator:
+    """JAX-traceable CosmoPower emulator.
+
+    Hold the static weights as jnp arrays; ``predict(params_dict)`` runs
+    the forward pass and returns the output (already de-standardised, and
+    PCA-decompressed if applicable).
+
+    Attributes
+    ----------
+    parameters : list[str]
+        Input parameter names, in training order.
+    modes : np.ndarray
+        Output mode labels (ell, k-index, etc. — interpretation is
+        emulator-specific). Length = output dimension AFTER PCA decompression.
+    is_log : bool
+        If True, the emulator output is log10 of the physical quantity.
+        Caller can apply ``10**`` if needed; we leave that to the caller
+        because the convention varies (e.g. TE doesn't apply log10 since it
+        can go negative).
+    """
+
+    # JAX-side state
+    weights:        tuple = field(repr=False)             # tuple of (W, b) per layer
+    alphas:         tuple = field(repr=False)             # α per hidden layer
+    betas:          tuple = field(repr=False)             # β per hidden layer
+    param_mean:     jnp.ndarray = field(repr=False)
+    param_std:      jnp.ndarray = field(repr=False)
+    feature_mean:   jnp.ndarray = field(repr=False)
+    feature_std:    jnp.ndarray = field(repr=False)
+    pca_matrix:     jnp.ndarray | None = field(repr=False, default=None)
+    pca_mean:       jnp.ndarray | None = field(repr=False, default=None)
+    pca_std:        jnp.ndarray | None = field(repr=False, default=None)
+
+    # Metadata
+    parameters:     tuple[str, ...] = ()
+    modes:          np.ndarray = field(default=None, repr=False)
+
+    @classmethod
+    def from_npz(cls, path: str) -> "Emulator":
+        """Load an emulator from a CosmoPower-style ``.npz`` file."""
+        raw = np.load(path, allow_pickle=True)
+        # CosmoPower saves a single object array containing a dict-like
+        d = raw["arr_0"].item()
+
+        # weights_ and biases_ are SEPARATE lists (one entry per layer);
+        # pair them up so the forward pass can iterate over (W, b) tuples.
+        weights = tuple(
+            (jnp.asarray(W, dtype=jnp.float64),
+             jnp.asarray(b, dtype=jnp.float64))
+            for W, b in zip(d["weights_"], d["biases_"])
+        )
+        # alphas_ / betas_: per hidden layer (one entry per hidden layer,
+        # i.e. len = n_layers - 1; final linear layer has no α-β activation)
+        alphas = tuple(jnp.asarray(a, dtype=jnp.float64) for a in d["alphas_"])
+        betas  = tuple(jnp.asarray(b, dtype=jnp.float64) for b in d["betas_"])
+
+        param_mean   = jnp.asarray(d["param_train_mean"],   dtype=jnp.float64)
+        param_std    = jnp.asarray(d["param_train_std"],    dtype=jnp.float64)
+        feature_mean = jnp.asarray(d["feature_train_mean"], dtype=jnp.float64)
+        feature_std  = jnp.asarray(d["feature_train_std"],  dtype=jnp.float64)
+
+        # PCA fields are present only for PCA emulators (e.g. TE)
+        has_pca = "pca_transform_matrix" in d
+        pca_matrix = jnp.asarray(d["pca_transform_matrix"], dtype=jnp.float64) if has_pca else None
+        pca_mean   = jnp.asarray(d["pca_mean"],   dtype=jnp.float64) if has_pca else None
+        pca_std    = jnp.asarray(d["pca_std"],    dtype=jnp.float64) if has_pca else None
+
+        parameters = tuple(d["parameters"])
+        modes      = np.asarray(d["modes"])
+
+        return cls(
+            weights=weights, alphas=alphas, betas=betas,
+            param_mean=param_mean, param_std=param_std,
+            feature_mean=feature_mean, feature_std=feature_std,
+            pca_matrix=pca_matrix, pca_mean=pca_mean, pca_std=pca_std,
+            parameters=parameters, modes=modes,
+        )
+
+    def predict(self, params_dict: dict) -> jnp.ndarray:
+        """Forward pass on a single or batched input.
+
+        ``params_dict``: maps each ``self.parameters`` name to a 1-D array-like
+        of length B (batch size); the output has shape ``(B, n_modes)`` (or
+        ``(n_modes,)`` if B=1).
+        """
+        # Stack inputs in the canonical order
+        p = jnp.stack([jnp.asarray(params_dict[name], dtype=jnp.float64)
+                       for name in self.parameters], axis=-1)   # (B, n_params)
+        # Standardise
+        h = (p - self.param_mean) / self.param_std
+
+        # Hidden layers with α-β sigmoid activation.
+        # CosmoPower .npz weights are stored as (in_dim, out_dim) so the
+        # matmul is `h @ W` (no transpose).
+        for (W, b), a, beta in zip(self.weights[:-1], self.alphas, self.betas):
+            z = h @ W + b
+            h = (beta + _sigmoid(a * z) * (1.0 - beta)) * z
+
+        # Final linear layer (no activation)
+        W, b = self.weights[-1]
+        preds = h @ W + b                                            # (B, n_out)
+
+        # PCA decompression (if applicable)
+        if self.pca_matrix is not None:
+            # PCA-encoded: undo PCA standardisation, project back, then
+            # undo feature standardisation
+            preds = preds * self.pca_std + self.pca_mean             # (B, n_pca)
+            preds = preds @ self.pca_matrix                          # (B, n_modes)
+
+        # Undo feature standardisation
+        preds = preds * self.feature_std + self.feature_mean
+        return preds.squeeze()
+
+
+# ---------------------------------------------------------------------------
+# Default data-directory resolution
+# ---------------------------------------------------------------------------
+
+def default_data_dir() -> str:
+    """Resolve the CosmoPower emulator data directory.
+
+    Resolution order:
+      1. ``$CLASSY_SZLITE_DATA_DIR`` environment variable
+      2. ``~/class_sz_data`` (the recommended new location)
+      3. ``~/class_sz_data_directory`` (legacy classy_szfast path)
+    """
+    env = os.environ.get("CLASSY_SZLITE_DATA_DIR")
+    if env:
+        return os.path.expanduser(env)
+    for candidate in ("~/class_sz_data", "~/class_sz_data_directory"):
+        path = os.path.expanduser(candidate)
+        if os.path.isdir(path):
+            return path
+    raise RuntimeError(
+        "classy_szlite: cannot find CosmoPower emulator data. Set the "
+        "$CLASSY_SZLITE_DATA_DIR environment variable, or place the emulator "
+        "directory tree at ~/class_sz_data (or the legacy "
+        "~/class_sz_data_directory). Required subdirs per cosmo_model: "
+        "PK/, growth-and-distances/, TTTEEE/, derived-parameters/."
+    )

classy_szlite/_registry.py

@@ -0,0 +1,94 @@
+"""Emulator file IDs, grid configs, and lazy loading.
+
+Single emulator suite: **ede-v2**. The ede-v2 CosmoPower emulators are
+extremely accurate and cover the LCDM-equivalent point at ``fEDE = 0.001``;
+they replace any need for a separate LCDM emulator path.
+"""
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+
+from ._emulator import Emulator, default_data_dir
+
+
+# File IDs (paths relative to <data_root>/ede/)
+_EMULATOR_FILES = {
+    "tt":   "TTTEEE/TT_v2.npz",
+    "te":   "TTTEEE/TE_v2.npz",
+    "ee":   "TTTEEE/EE_v2.npz",
+    "pp":   "PP/PP_v2.npz",
+    "pkl":  "PK/PKL_v2.npz",
+    "pknl": "PK/PKNL_v2.npz",
+    "hz":   "growth-and-distances/HZ_v2.npz",
+    "daz":  "growth-and-distances/DAZ_v2.npz",
+    "s8z":  "growth-and-distances/S8Z_v2.npz",
+    "der":  "derived-parameters/DER_v2.npz",
+}
+
+# Subdir under the data root
+_DATA_SUBDIR = "ede"
+
+
+# ---------------------------------------------------------------------------
+# Pk emulator grid + extrapolation configuration
+# ---------------------------------------------------------------------------
+# Final n_k = nk // ndspl. Final k = geomspace(kmin, kmax, nk)[::ndspl].
+# prefac:
+#   'k3' — emulator output is log10[ k³ Pk ]; recover Pk via × k^-3
+# extrap_kmin: extend the Pk grid down to this k via P(k) ∝ k^n_s
+PK_GRID = dict(kmin=5e-4, kmax=10.0, nk=1000, ndspl=1,
+               prefac="k3", extrap_kmin=1e-4)
+
+# Distance emulator z-grid: HZ trained on linspace(0, 20, 5000); DAZ omits
+# z=0 (returns 4999 values) — we prepend chi(0)=0 in cosmology._predict_distances.
+DIST_ZMAX = 20.0
+
+
+# ---------------------------------------------------------------------------
+# Per-emulator log10 convention
+# ---------------------------------------------------------------------------
+# Whether the raw NN output is log10(physical) vs physical directly.
+# True  → caller must apply 10^ to recover the physical quantity.
+# False → raw output IS the physical quantity (or the prefactored-Pk form).
+LOG10_OUTPUT = {
+    "tt": True, "ee": True, "pp": True, "te": False,    # CMB
+    "pkl": True, "pknl": True,                          # Pk (handled in cosmology)
+    "hz": True, "daz": True, "s8z": True,               # distances + sigma8(z)
+    "der": True,                                        # derived (sigma_8 = der[1])
+}
+
+
+# ---------------------------------------------------------------------------
+# Default cosmology parameters (matches ede-v2 emulator training default)
+# ---------------------------------------------------------------------------
+DEFAULT_COSMO = {
+    "ln10^{10}A_s": 3.047, "n_s": 0.9665,
+    "H0": 67.66, "omega_b": 0.02242, "omega_cdm": 0.11933,
+    "tau_reio": 0.054,
+    # EDE-specific defaults (LCDM-equivalent point)
+    "fEDE": 0.001, "log10z_c": 3.562, "thetai_scf": 2.83, "r": 0.0,
+    # ν convention: 3 degenerate ncdm of 0.02 eV → Σmν = 0.06 eV
+    "m_ncdm": 0.02, "N_ur": 0.00441,
+}
+
+
+# ---------------------------------------------------------------------------
+# Lazy emulator loading + cache
+# ---------------------------------------------------------------------------
+
+@lru_cache(maxsize=16)
+def get_emulator(name: str) -> Emulator:
+    """Load (and cache) the named emulator."""
+    if name not in _EMULATOR_FILES:
+        raise ValueError(f"Unknown emulator name {name!r}. "
+                         f"Supported: {list(_EMULATOR_FILES)}")
+    root = default_data_dir()
+    path = os.path.join(root, _DATA_SUBDIR, _EMULATOR_FILES[name])
+    if not os.path.isfile(path):
+        raise FileNotFoundError(f"classy_szlite: emulator file not found: {path}")
+    return Emulator.from_npz(path)
+
+
+def output_is_log10(name: str) -> bool:
+    return LOG10_OUTPUT[name]

classy_szlite/api.py

@@ -0,0 +1,194 @@
+"""High-level public API: derived params, CMB Cls, matter Pk, distances,
+tSZ Cl^yy. All functions take a :class:`classy_szlite.params.CosmoParams`.
+"""
+from __future__ import annotations
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+from ._registry import get_emulator, DEFAULT_COSMO
+from .params import CosmoParams, ProfileParamsA10
+from .cosmology import build as build_cosmo_grids, get_pk, get_pknl, get_distances
+from .hmf import build_halo_grids
+from .power_spectrum import cl_yy_1h_2h
+
+jax.config.update("jax_enable_x64", True)
+
+
+def cosmo_to_dict(cosmo: CosmoParams) -> dict:
+    """Convert ``CosmoParams`` to the emulator-style dict (with curly-brace key).
+
+    Values pass through as-is (float or jax.Array) so the returned dict is
+    JAX-traceable — pass a ``CosmoParams`` of tracers to ``jax.grad``.
+    """
+    return {
+        "omega_b":       cosmo.omega_b,
+        "omega_cdm":     cosmo.omega_cdm,
+        "H0":            cosmo.H0,
+        "tau_reio":      cosmo.tau_reio,
+        "ln10^{10}A_s":  cosmo.ln10_10_As,
+        "n_s":           cosmo.n_s,
+        "m_ncdm":        cosmo.m_ncdm,
+        "N_ur":          cosmo.N_ur,
+        "fEDE":          cosmo.fEDE,
+        "log10z_c":      cosmo.log10z_c,
+        "thetai_scf":    cosmo.thetai_scf,
+        "r":             cosmo.r,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Derived parameters (σ8, Ω_m, S8)
+# ---------------------------------------------------------------------------
+
+def derived(cosmo: CosmoParams) -> dict:
+    """Derived parameters: ``sigma_8``, ``Omega_m``, ``S8``.
+
+    Also returns the full 17-element DER emulator output as ``'der_full'``
+    (σ8 is at index 1; consult the CosmoPower DER training script for the
+    full list).
+    """
+    em = get_emulator('der')
+    full = dict(DEFAULT_COSMO); full.update(cosmo_to_dict(cosmo))
+    p_in = {k: [full[k]] for k in em.parameters}
+    der = 10.0 ** np.asarray(em.predict(p_in)).flatten()
+
+    sigma_8 = float(der[1])
+    h = float(cosmo.H0) / 100.0
+    # Σmν = 3 × m_ncdm for ede-v2 ν convention (3 degenerate ν)
+    sum_mnu = 3.0 * float(cosmo.m_ncdm)
+    Omega_m = (float(cosmo.omega_b) + float(cosmo.omega_cdm)
+               + sum_mnu / 93.14) / h ** 2
+    S8 = sigma_8 * (Omega_m / 0.3) ** 0.5
+    return {"sigma_8": sigma_8, "Omega_m": Omega_m, "S8": S8, "der_full": der}
+
+
+# ---------------------------------------------------------------------------
+# CMB angular power spectra
+# ---------------------------------------------------------------------------
+
+# TT, EE, PP emulators output log10(prefactored Cl); TE outputs Cl directly
+# (it can be negative). Recovery factor: ede-v2 uses 1/ell² to get raw Cl.
+_CMB_LOG_CONVENTION = {"tt": True, "ee": True, "pp": True, "te": False}
+
+
+def cl_TTTEEE(cosmo: CosmoParams,
+              spectra: tuple[str, ...] = ("tt", "te", "ee"),
+              ell_factor: bool = True) -> dict:
+    """CMB angular power spectra.
+
+    Returns a dict with keys ``'ell'`` and the requested spectra
+    (``'tt','te','ee'``). Values are **dimensionless** — multiply by
+    ``Tcmb_uK² = (2.7255e6)²`` to convert to μK².
+
+    ``ell_factor`` (default ``True``) — return ``D_ell = ell(ell+1) Cl / (2π)``;
+    ``False`` returns raw Cl.
+    """
+    full = dict(DEFAULT_COSMO); full.update(cosmo_to_dict(cosmo))
+
+    out = {}
+    ell = None
+    for spec in spectra:
+        if spec not in _CMB_LOG_CONVENTION:
+            raise ValueError(f"Unknown spectrum {spec!r}. Pick from {tuple(_CMB_LOG_CONVENTION)}.")
+        em = get_emulator(spec)
+        p_in = {k: [full[k]] for k in em.parameters}
+        pred = np.asarray(em.predict(p_in)).flatten()
+        if _CMB_LOG_CONVENTION[spec]:
+            pred = 10.0 ** pred
+        out[spec] = pred
+        if ell is None:
+            ell = np.asarray(em.modes)
+
+    # ede-v2 recovery factor: raw → Cl
+    factor_to_Cl = 1.0 / (ell ** 2)
+    for s in spectra:
+        out[s] = out[s] * factor_to_Cl
+
+    out["ell"] = ell
+    if ell_factor:
+        fac_dl = ell * (ell + 1) / (2.0 * np.pi)
+        for s in spectra:
+            out[s] = out[s] * fac_dl
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Pk, Pnl, distances
+# ---------------------------------------------------------------------------
+
+def Pk(cosmo: CosmoParams, z_arr):
+    """Linear P(k, z) — returns ``(k, pk(z, k))``."""
+    return get_pk(cosmo_to_dict(cosmo), z_arr)
+
+
+def Pnl(cosmo: CosmoParams, z_arr):
+    """Non-linear P(k, z) (HMcode) — returns ``(k, pk(z, k))``."""
+    return get_pknl(cosmo_to_dict(cosmo), z_arr)
+
+
+def distances(cosmo: CosmoParams, z_arr):
+    """Returns ``(Hz, chi, Da)``. ``Hz`` is H(z)/c in 1/Mpc; distances in Mpc."""
+    return get_distances(cosmo_to_dict(cosmo), z_arr)
+
+
+# ---------------------------------------------------------------------------
+# tSZ Cl^yy (halo-model, Arnaud 2010 profile)
+# ---------------------------------------------------------------------------
+
+def cl_yy(cosmo: CosmoParams, profile: ProfileParamsA10, ell,
+          z_grid: jax.Array | None = None,
+          n_z: int = 100, m_min: float = 1e10, m_max: float = 3.5e15,
+          n_m: int = 200, delta_crit: float = 500.0):
+    """Halo-model tSZ angular power spectrum (full pipeline per call).
+
+    Returns ``(cl_1h, cl_2h)`` — dimensionless C_ell. Multiply by
+    ``ell*(ell+1)/(2π)*1e12`` to get ``D_ell × 1e12``.
+
+    For MCMC sampling only profile parameters at fixed cosmology, use
+    :func:`cl_yy_factory` instead — ~3× faster.
+    """
+    if z_grid is None:
+        z_grid = jnp.geomspace(0.005, 3.0, n_z)
+    cosmo_dict = cosmo_to_dict(cosmo)
+    cg = build_cosmo_grids(cosmo_dict, z_grid=z_grid)
+    hg = build_halo_grids(cg, cosmo_dict, delta_crit=delta_crit,
+                          m_min=m_min, m_max=m_max, n_m=n_m)
+    pp_dict = profile._asdict()
+    cl_1h, cl_2h = cl_yy_1h_2h(jnp.asarray(ell), cg, hg, cosmo_dict,
+                                profile='arnaud10', profile_params=pp_dict)
+    return cl_1h, cl_2h
+
+
+def cl_yy_factory(cosmo: CosmoParams, ell,
+                  z_grid: jax.Array | None = None,
+                  n_z: int = 100, m_min: float = 1e10, m_max: float = 3.5e15,
+                  n_m: int = 200, delta_crit: float = 500.0):
+    """Fixed-cosmology fast-path: precompute the heavy bits, get a closure.
+
+    Builds ``CosmoGrids`` (emulators → P_lin, distances, σ(R)) and
+    ``HaloGrids`` (Tinker 08 HMF, bias) **once**, then returns:
+
+        ev(profile) -> (cl_1h, cl_2h)
+
+    A subsequent ``ev(profile)`` call only runs the ``cl_yy_1h_2h``
+    halo-model integration — typically ~5 ms per call. Intended for MCMC
+    over profile / nuisance parameters with fixed cosmology.
+    """
+    if z_grid is None:
+        z_grid = jnp.geomspace(0.005, 3.0, n_z)
+    cosmo_dict = cosmo_to_dict(cosmo)
+    cg = build_cosmo_grids(cosmo_dict, z_grid=z_grid)
+    hg = build_halo_grids(cg, cosmo_dict, delta_crit=delta_crit,
+                          m_min=m_min, m_max=m_max, n_m=n_m)
+    ell_jax = jnp.asarray(ell)
+
+    def evaluate(profile: ProfileParamsA10):
+        cl_1h, cl_2h = cl_yy_1h_2h(
+            ell_jax, cg, hg, cosmo_dict,
+            profile='arnaud10', profile_params=profile._asdict(),
+        )
+        return cl_1h, cl_2h
+
+    return evaluate