classy-szlite 0.2.0__tar.gz

classy_szlite-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Boris Bolliet and the CLASS-SZ contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

classy_szlite-0.2.0/PKG-INFO

@@ -0,0 +1,148 @@
+Metadata-Version: 2.4
+Name: classy_szlite
+Version: 0.2.0
+Summary: Fast, differentiable cosmology in pure JAX. CMB Cls, matter Pk (linear and nonlinear), distances, derived parameters, and halo-model tSZ Cl^yy via high-accuracy ede-v2 CosmoPower emulators.
+Author-email: Boris Bolliet <boris.bolliet@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/CLASS-SZ/classy_szlite
+Project-URL: Source, https://github.com/CLASS-SZ/classy_szlite
+Keywords: cosmology,cmb,tsz,halo-model,jax,ede
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4
+Requires-Dist: numpy>=1.24
+Requires-Dist: mcfit>=0.0.21
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: pytest-xdist; extra == "dev"
+Provides-Extra: plots
+Requires-Dist: matplotlib>=3.6; extra == "plots"
+Dynamic: license-file
+
+# classy_szlite
+
+[![tests](https://github.com/CLASS-SZ/classy_szlite/actions/workflows/test.yml/badge.svg)](https://github.com/CLASS-SZ/classy_szlite/actions/workflows/test.yml)
+[![PyPI](https://img.shields.io/pypi/v/classy_szlite.svg)](https://pypi.org/project/classy_szlite/)
+[![Documentation Status](https://readthedocs.org/projects/classy-szlite/badge/?version=latest)](https://classy-szlite.readthedocs.io/en/latest/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+**Fast, differentiable cosmology in pure JAX.**
+
+`classy_szlite` provides JIT-compiled, `jax.grad`-friendly access to:
+
+- **CMB angular power spectra** — TT, TE, EE
+- **Linear and nonlinear matter Pk** — P(k, z), Pnl(k, z)
+- **Distances** — H(z), comoving χ(z), angular-diameter D_A(z)
+- **Derived parameters** — σ8, Ω_m, S8
+- **Halo-model tSZ Cl^yy** — Arnaud 2010 GNFW pressure profile
+
+Backed by the high-accuracy **ede-v2 CosmoPower emulators**. LCDM-equivalent
+cosmology is the default — `fEDE = 0.001` (the package's default) reproduces
+standard ΛCDM to emulator precision, and the same emulator covers the full
+early-dark-energy parameter space if you want to explore it.
+
+Runtime dependencies: `jax`, `numpy`, `mcfit`.
+
+## Install
+
+```bash
+pip install classy_szlite
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/CLASS-SZ/classy_szlite
+cd classy_szlite
+pip install -e .
+```
+
+You also need the CosmoPower emulator `.npz` files at `~/class_sz_data/`
+(or the path in `$CLASSY_SZLITE_DATA_DIR`). See
+[Installation](https://classy-szlite.readthedocs.io/en/latest/installation.html).
+
+## Quick start
+
+```python
+import jax.numpy as jnp
+import classy_szlite as csl
+
+cosmo = csl.CosmoParams()                      # Planck-18 + ede-v2 LCDM-equivalent defaults
+
+# Derived parameters
+csl.derived(cosmo)
+# → {'sigma_8': 0.812, 'Omega_m': 0.311, 'S8': 0.827, 'der_full': ...}
+
+# CMB Cls (dimensionless D_ℓ; × Tcmb² for μK²)
+csl.cl_TTTEEE(cosmo)
+# → {'ell', 'tt', 'te', 'ee'}
+
+# Matter Pk at multiple z
+k, pk  = csl.Pk(cosmo,  [0., 0.5, 1., 2.])
+k, pnl = csl.Pnl(cosmo, [0., 0.5, 1., 2.])
+
+# Distances
+Hz, chi, Da = csl.distances(cosmo, [0.1, 0.5, 1.0])
+
+# Halo-model tSZ Cl^yy
+profile = csl.ProfileParamsA10(P0=8.13, beta=5.48, B=1.25)
+ell = jnp.geomspace(2, 9000, 80)
+cl_1h, cl_2h = csl.cl_yy(cosmo, profile, ell)
+
+# MCMC fast path: precompute cosmology + halo grids once → ~5 ms/call
+ev = csl.cl_yy_factory(cosmo, ell)
+cl_1h, cl_2h = ev(profile)
+```
+
+## Throughput
+
+Warm-call timing, n = 100 calls per benchmark, freshly randomised inputs:
+
+| Function | mean ± std (ms) | calls/s |
+| --- | --- | --- |
+| `derived` | 0.54 ± 0.04 | 1850 ± 150 |
+| `cl_TTTEEE` | 2.52 ± 0.14 | 400 ± 25 |
+| `Pk` | 1.49 ± 0.12 | 670 ± 55 |
+| `distances` | 1.29 ± 0.09 | 770 ± 60 |
+| `cl_yy` (full pipeline) | 17.84 ± 0.58 | 56 ± 2 |
+| `cl_yy_factory` (fixed-cosmo) | 5.38 ± 0.42 | 185 ± 15 |
+| `cl_yy_factory + jax.grad` | 17.12 ± 1.01 | 58 ± 3 |
+
+Reference platform: macOS arm64 (M-series CPU), single-thread JAX. See
+[Throughput](https://classy-szlite.readthedocs.io/en/latest/throughput.html)
+for a more detailed table + reproduction script.
+
+## JAX gradients
+
+All public functions are JAX-traceable. The factory closure is the
+recommended path for gradient-based inference at fixed cosmology:
+
+```python
+import jax
+ev = csl.cl_yy_factory(cosmo, ell)
+
+def loss(P0, beta):
+    cl_1h, cl_2h = ev(csl.ProfileParamsA10(P0=P0, beta=beta, B=1.25))
+    return jnp.sum(cl_1h + cl_2h)
+
+d_loss = jax.grad(loss, argnums=(0, 1))(8.13, 5.48)
+```
+
+Gradients also work through the full pipeline (cosmology + profile) and
+w.r.t. the whole `CosmoParams` pytree — see
+[Gradients](https://classy-szlite.readthedocs.io/en/latest/gradients.html).
+
+## Documentation
+
+Full docs at **https://classy-szlite.readthedocs.io**.
+
+## License
+
+MIT.

classy_szlite-0.2.0/README.md

@@ -0,0 +1,120 @@
+# classy_szlite
+
+[![tests](https://github.com/CLASS-SZ/classy_szlite/actions/workflows/test.yml/badge.svg)](https://github.com/CLASS-SZ/classy_szlite/actions/workflows/test.yml)
+[![PyPI](https://img.shields.io/pypi/v/classy_szlite.svg)](https://pypi.org/project/classy_szlite/)
+[![Documentation Status](https://readthedocs.org/projects/classy-szlite/badge/?version=latest)](https://classy-szlite.readthedocs.io/en/latest/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+**Fast, differentiable cosmology in pure JAX.**
+
+`classy_szlite` provides JIT-compiled, `jax.grad`-friendly access to:
+
+- **CMB angular power spectra** — TT, TE, EE
+- **Linear and nonlinear matter Pk** — P(k, z), Pnl(k, z)
+- **Distances** — H(z), comoving χ(z), angular-diameter D_A(z)
+- **Derived parameters** — σ8, Ω_m, S8
+- **Halo-model tSZ Cl^yy** — Arnaud 2010 GNFW pressure profile
+
+Backed by the high-accuracy **ede-v2 CosmoPower emulators**. LCDM-equivalent
+cosmology is the default — `fEDE = 0.001` (the package's default) reproduces
+standard ΛCDM to emulator precision, and the same emulator covers the full
+early-dark-energy parameter space if you want to explore it.
+
+Runtime dependencies: `jax`, `numpy`, `mcfit`.
+
+## Install
+
+```bash
+pip install classy_szlite
+```
+
+Or from source:
+
+```bash
+git clone https://github.com/CLASS-SZ/classy_szlite
+cd classy_szlite
+pip install -e .
+```
+
+You also need the CosmoPower emulator `.npz` files at `~/class_sz_data/`
+(or the path in `$CLASSY_SZLITE_DATA_DIR`). See
+[Installation](https://classy-szlite.readthedocs.io/en/latest/installation.html).
+
+## Quick start
+
+```python
+import jax.numpy as jnp
+import classy_szlite as csl
+
+cosmo = csl.CosmoParams()                      # Planck-18 + ede-v2 LCDM-equivalent defaults
+
+# Derived parameters
+csl.derived(cosmo)
+# → {'sigma_8': 0.812, 'Omega_m': 0.311, 'S8': 0.827, 'der_full': ...}
+
+# CMB Cls (dimensionless D_ℓ; × Tcmb² for μK²)
+csl.cl_TTTEEE(cosmo)
+# → {'ell', 'tt', 'te', 'ee'}
+
+# Matter Pk at multiple z
+k, pk  = csl.Pk(cosmo,  [0., 0.5, 1., 2.])
+k, pnl = csl.Pnl(cosmo, [0., 0.5, 1., 2.])
+
+# Distances
+Hz, chi, Da = csl.distances(cosmo, [0.1, 0.5, 1.0])
+
+# Halo-model tSZ Cl^yy
+profile = csl.ProfileParamsA10(P0=8.13, beta=5.48, B=1.25)
+ell = jnp.geomspace(2, 9000, 80)
+cl_1h, cl_2h = csl.cl_yy(cosmo, profile, ell)
+
+# MCMC fast path: precompute cosmology + halo grids once → ~5 ms/call
+ev = csl.cl_yy_factory(cosmo, ell)
+cl_1h, cl_2h = ev(profile)
+```
+
+## Throughput
+
+Warm-call timing, n = 100 calls per benchmark, freshly randomised inputs:
+
+| Function | mean ± std (ms) | calls/s |
+| --- | --- | --- |
+| `derived` | 0.54 ± 0.04 | 1850 ± 150 |
+| `cl_TTTEEE` | 2.52 ± 0.14 | 400 ± 25 |
+| `Pk` | 1.49 ± 0.12 | 670 ± 55 |
+| `distances` | 1.29 ± 0.09 | 770 ± 60 |
+| `cl_yy` (full pipeline) | 17.84 ± 0.58 | 56 ± 2 |
+| `cl_yy_factory` (fixed-cosmo) | 5.38 ± 0.42 | 185 ± 15 |
+| `cl_yy_factory + jax.grad` | 17.12 ± 1.01 | 58 ± 3 |
+
+Reference platform: macOS arm64 (M-series CPU), single-thread JAX. See
+[Throughput](https://classy-szlite.readthedocs.io/en/latest/throughput.html)
+for a more detailed table + reproduction script.
+
+## JAX gradients
+
+All public functions are JAX-traceable. The factory closure is the
+recommended path for gradient-based inference at fixed cosmology:
+
+```python
+import jax
+ev = csl.cl_yy_factory(cosmo, ell)
+
+def loss(P0, beta):
+    cl_1h, cl_2h = ev(csl.ProfileParamsA10(P0=P0, beta=beta, B=1.25))
+    return jnp.sum(cl_1h + cl_2h)
+
+d_loss = jax.grad(loss, argnums=(0, 1))(8.13, 5.48)
+```
+
+Gradients also work through the full pipeline (cosmology + profile) and
+w.r.t. the whole `CosmoParams` pytree — see
+[Gradients](https://classy-szlite.readthedocs.io/en/latest/gradients.html).
+
+## Documentation
+
+Full docs at **https://classy-szlite.readthedocs.io**.
+
+## License
+
+MIT.

classy_szlite-0.2.0/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-0.2.0/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-0.2.0/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]