Cosmological Models

HIcosmo provides multiple cosmological models, covering the standard cosmological model and its dynamic dark energy extensions. All models are built on JAX with support for JIT compilation and automatic differentiation.

Model Overview

Model	Dark Energy Equation of State	Description
LCDM	\(w = -1\)	Standard \(\Lambda\text{CDM}\) model, cosmological constant
wCDM	\(w = w_0\)	Constant dark energy equation of state
CPL	\(w(a) = w_0 + w_a(1-a)\)	Chevallier-Polarski-Linder parameterization
ILCDM	\(w = -1\) + interaction	Interacting dark energy model

Basic Theory

Dimensionless Hubble Parameter

The core of all models is the dimensionless Hubble parameter \(E(z) = H(z)/H_0\). It determines the expansion history of the universe and all distance measurements.

General form:

\[E^2(z) = \Omega_m (1+z)^3 + \Omega_r (1+z)^4 + \Omega_k (1+z)^2 + \Omega_{\rm DE}(z)\]

where:

• \(\Omega_m\): Matter density parameter

• \(\Omega_r\): Radiation density parameter

• \(\Omega_k\): Curvature density parameter

• \(\Omega_{\rm DE}(z)\): Dark energy density evolution

Normalization condition: At \(z=0\), \(E(0) = 1\), which requires:

\[\Omega_m + \Omega_r + \Omega_k + \Omega_\Lambda = 1\]

Distance Calculations

All cosmological distances can be derived from \(E(z)\):

Comoving distance:

\[d_C(z) = \frac{c}{H_0} \int_0^z \frac{dz'}{E(z')}\]

Transverse comoving distance (accounting for curvature):

\[\begin{split}D_M(z) = \begin{cases} D_H / \sqrt{\Omega_k} \sinh(\sqrt{\Omega_k} \, d_C / D_H) & \Omega_k > 0 \\ d_C & \Omega_k = 0 \\ D_H / \sqrt{|\Omega_k|} \sin(\sqrt{|\Omega_k|} \, d_C / D_H) & \Omega_k < 0 \end{cases}\end{split}\]

where \(D_H = c/H_0\) is the Hubble distance.

Luminosity distance:

\[d_L(z) = (1+z) \, D_M(z)\]

Angular diameter distance:

\[D_A(z) = \frac{D_M(z)}{1+z}\]

Distance modulus:

\[\mu = 5 \log_{10}\left(\frac{d_L}{\text{Mpc}}\right) + 25\]

LCDM Model

The standard \(\Lambda\text{CDM}\) model assumes dark energy is a cosmological constant, \(w = -1\).

LCDM Mathematical Form

Hubble parameter:

\[E^2(z) = \Omega_m (1+z)^3 + \Omega_r (1+z)^4 + \Omega_k (1+z)^2 + \Omega_\Lambda\]

Dark energy equation of state:

\[w(z) = -1 \quad \text{(constant)}\]

LCDM Parameters

Parameter	Default	Range	Description
H0	67.36	[50, 100]	Hubble constant in km/s/Mpc
Omega_m	0.3153	[0.01, 0.99]	Total matter density parameter
Omega_b	0.0493	[0.01, 0.1]	Baryon density parameter
Omega_k	0.0	[-0.3, 0.3]	Curvature density parameter (0 = flat universe)
Omega_r	Auto-computed	–	Radiation density parameter, computed from \(T_{\rm CMB}\)
sigma8	0.8111	[0.5, 1.2]	Matter power spectrum normalization
n_s	0.9649	[0.8, 1.1]	Primordial power spectrum scalar spectral index
T_cmb	2.7255	–	CMB temperature (K)
N_eff	3.046	–	Effective number of neutrino species

LCDM Basic Usage

from hicosmo.models import LCDM

# Use Planck 2018 default parameters
model = LCDM()

# Or specify parameters
model = LCDM(H0=70.0, Omega_m=0.3, Omega_k=0.01)

# Compute distances
z = 1.0
d_L = model.luminosity_distance(z)      # Luminosity distance [Mpc]
D_A = model.angular_diameter_distance(z)  # Angular diameter distance [Mpc]
D_M = model.transverse_comoving_distance(z)  # Transverse comoving distance [Mpc]
mu = model.distance_modulus(z)          # Distance modulus

print(f"d_L(z=1) = {d_L:.2f} Mpc")
print(f"D_A(z=1) = {D_A:.2f} Mpc")
print(f"mu(z=1) = {mu:.2f}")

Parameter Consistency and Updates

All model parameters use model.params as the single source of truth (internally CosmologicalParameters). Model attributes w / w0 / wa / beta are lightweight wrappers around params, ensuring consistency between instance methods and compute_grid_traced physics calculations.

Recommended way to update parameters:

from hicosmo.models import LCDM, wCDM, CPL, ILCDM

model = LCDM()
model.params['H0'] = 70.0
model.params['Omega_m'] = 0.3

w_model = wCDM(w=-1.0)
w_model.w = -0.9                 # Equivalent to w_model.params['w'] = -0.9

cpl = CPL(w0=-1.0, wa=0.0)
cpl.w0 = -0.95                   # Equivalent to cpl.params['w0'] = -0.95
cpl.wa = 0.2

ilcdm = ILCDM(beta=0.0)
ilcdm.beta = 0.05                # Equivalent to ilcdm.params['beta'] = 0.05

Sound Horizon Calculation

The LCDM model includes the calculation of the sound horizon at the drag epoch \(r_d\), a key quantity for BAO analysis.

Sound horizon (Eisenstein & Hu 1998 fitting formula):

\[r_d = \frac{44.5 \ln(9.83 / \Omega_m h^2)}{\sqrt{1 + 10 (\Omega_b h^2)^{3/4}}} \text{ Mpc}\]

# Get the sound horizon
rd = model.sound_horizon()  # Mpc
rd_h = model.rd_h           # r_d * h (dimensionless)

print(f"r_d = {rd:.2f} Mpc")

Growth Functions

The LCDM model provides calculations of the linear growth factor and growth rate.

Growth factor D(z) (Carroll, Press & Turner 1992 approximation):

\[D(z) \propto \frac{5 \Omega_m E(z)}{2} \int_z^\infty \frac{(1+z') dz'}{E^3(z')}\]

Growth rate:

\[f(z) = \frac{d \ln D}{d \ln a} \approx \Omega_m(z)^\gamma, \quad \gamma \approx 0.55\]

# Growth factor (normalized to z=0)
D = model.growth_factor(z=0.5)

# Growth rate
f = model.growth_rate(z=0.5)

# f*sigma8 parameter (commonly used in RSD analysis)
fsig8 = model.fsigma8(z=0.5)

print(f"D(z=0.5) = {D:.4f}")
print(f"f(z=0.5) = {f:.4f}")
print(f"f*sigma8(z=0.5) = {fsig8:.4f}")

wCDM Model

The wCDM model generalizes the dark energy equation of state to a constant \(w \neq -1\).

wCDM Mathematical Form

Hubble parameter:

\[E^2(z) = \Omega_m (1+z)^3 + \Omega_r (1+z)^4 + \Omega_k (1+z)^2 + \Omega_\Lambda (1+z)^{3(1+w)}\]

Dark energy equation of state:

\[w(z) = w_0 \quad \text{(constant)}\]

wCDM Parameters

wCDM inherits all LCDM parameters and adds:

Parameter	Default	Prior Range	Description
w (or w0)	-1.0	[-2.5, 0.0]	Constant dark energy equation of state

Physical constraints:

• \(w = -1\): Cosmological constant (LCDM)

• \(w > -1\): Quintessence-type dark energy

• \(w < -1\): Phantom dark energy (violates energy conditions)

wCDM Basic Usage

from hicosmo.models import wCDM

# Create a wCDM model
model = wCDM(H0=70.0, Omega_m=0.3, w=-0.9)

# All distance methods inherited from LCDM
d_L = model.luminosity_distance(1.0)

# View the equation of state
w = model.w_z(0.5)  # Returns constant w
print(f"w(z=0.5) = {w:.2f}")

# E(z) now includes the w parameter
E = model.E_z(1.0)
print(f"E(z=1) = {E:.4f}")

CPL Model

The CPL (Chevallier-Polarski-Linder) parameterization allows the dark energy equation of state to evolve with time.

CPL Mathematical Form

Equation of state parameterization:

\[w(a) = w_0 + w_a (1 - a) = w_0 + w_a \frac{z}{1+z}\]

where \(a = 1/(1+z)\) is the scale factor.

Dark energy density evolution:

\[f_{\rm DE}(z) = \exp\left( \frac{3 w_a z}{1+z} \right) \cdot (1+z)^{3(1 + w_0 + w_a)}\]

Hubble parameter:

\[E^2(z) = \Omega_m (1+z)^3 + \Omega_r (1+z)^4 + \Omega_k (1+z)^2 + \Omega_\Lambda \, f_{\rm DE}(z)\]

CPL Parameters

CPL inherits all LCDM parameters and adds:

Parameter	Default	Prior Range	Description
w0	-1.0	[-2.5, 0.5]	Dark energy equation of state today (\(z=0\))
wa	0.0	[-3.0, 2.0]	Equation of state evolution parameter

Special cases:

• \(w_0 = -1, w_a = 0\): LCDM

• \(w_a = 0\): wCDM (constant \(w\))

• \(w_a \neq 0\): Dynamic dark energy

CPL Basic Usage

from hicosmo.models import CPL

# Create a CPL model
model = CPL(H0=67.36, Omega_m=0.3, w0=-1.0, wa=0.1)

# View the evolving equation of state
z_values = [0.0, 0.5, 1.0, 2.0]
for z in z_values:
    w = model.w_z(z)
    print(f"w(z={z}) = {w:.3f}")

# Output:
# w(z=0.0) = -1.000
# w(z=0.5) = -0.967
# w(z=1.0) = -0.950
# w(z=2.0) = -0.933

# Distance calculation
d_L = model.luminosity_distance(1.0)
print(f"d_L(z=1) = {d_L:.2f} Mpc")

Dark Energy Figure of Merit

The CPL parameterization is commonly used to compute the dark energy Figure of Merit (FoM):

\[\text{FoM} = \frac{1}{\sqrt{\det(\text{Cov}(w_0, w_a))}}\]

A larger FoM indicates tighter constraints on the dark energy equation of state.

ILCDM Model

The ILCDM (Interacting Lambda-CDM) model assumes energy exchange between dark matter and dark energy.

Physical Picture

In standard LCDM, dark matter and dark energy evolve independently. ILCDM introduces an interaction term:

\[Q = \beta H \rho_c\]

where \(\beta\) is a dimensionless coupling constant and \(\rho_c\) is the cold dark matter density.

Energy conservation equations:

\[\dot{\rho}_c + 3H\rho_c = -Q = -\beta H \rho_c\]

\[\dot{\rho}_\Lambda = +Q = \beta H \rho_c\]

Density evolution:

\[\rho_c(z) = \rho_{c,0} (1+z)^{3 - \beta}\]

ILCDM Mathematical Form

Hubble parameter:

\[E^2(z) = \Omega_\Lambda + \Omega_b (1+z)^3 + \Omega_r (1+z)^4 + \Omega_c \left[ \frac{\beta}{\beta - 1} + \frac{1}{1-\beta} (1+z)^{3-3\beta} \right]\]

where \(\Omega_c = \Omega_m - \Omega_b\) is the cold dark matter density parameter.

ILCDM Parameters

Parameter	Default	Prior Range	Description
beta	0.0	[-1.0, 1.0]	Dark energy–dark matter coupling constant

Physical meaning:

• \(\beta = 0\): Standard LCDM

• \(\beta > 0\): Dark matter decays into dark energy

• \(\beta < 0\): Dark energy decays into dark matter

Note: ILCDM assumes a flat universe (\(\Omega_k = 0\)).

ILCDM Basic Usage

from hicosmo.models import ILCDM

# Create an ILCDM model
model = ILCDM(H0=67.36, Omega_m=0.3, beta=0.05)

# Distance calculation
d_L = model.luminosity_distance(1.0)

# Compare different beta values
for beta in [-0.1, 0.0, 0.1]:
    m = ILCDM(H0=67.36, Omega_m=0.3, beta=beta)
    d_L = m.luminosity_distance(1.0)
    print(f"beta={beta:+.1f}: d_L(z=1) = {d_L:.2f} Mpc")

Traced-Aware Interface

For compatibility with the NumPyro MCMC sampler, all models provide a compute_grid_traced static method that supports JAX tracer propagation.

Basic Usage

from hicosmo.models import LCDM, CPL
import jax.numpy as jnp

# Define a redshift grid
z_grid = jnp.linspace(0.01, 2.0, 1000)

# Parameter dictionary (for MCMC sampling)
# Only H0 / Omega_m and model-specific parameters are needed
params = {'H0': 70.0, 'Omega_m': 0.3}

# Batch-compute all distances
grid = LCDM.compute_grid_traced(z_grid, params)

# The returned dictionary contains multiple quantities
d_L = grid['d_L']     # Luminosity distance
D_M = grid['D_M']     # Transverse comoving distance
D_H = grid['D_H']     # Hubble distance (c/H(z))
E_z = grid['E_z']     # Dimensionless Hubble parameter
d_C = grid['d_C']     # Comoving distance
dVc_dz = grid['dVc_dz']  # Comoving volume element

CPL Model Example

# CPL requires additional parameters
cpl_params = {
    'H0': 67.36,
    'Omega_m': 0.3,
    'w0': -1.0,
    'wa': 0.1
}

grid = CPL.compute_grid_traced(z_grid, cpl_params)
d_L = grid['d_L']

wCDM / ILCDM Parameter Dictionary Examples

from hicosmo.models import wCDM, ILCDM

w_params = {'H0': 70.0, 'Omega_m': 0.3, 'w': -0.9}
grid_w = wCDM.compute_grid_traced(z_grid, w_params)

ilcdm_params = {'H0': 70.0, 'Omega_m': 0.3, 'beta': 0.05}
grid_ilcdm = ILCDM.compute_grid_traced(z_grid, ilcdm_params)

Using with NumPyro

import numpyro
import numpyro.distributions as dist

def model(z_obs, mu_obs, mu_err):
    # Sample cosmological parameters
    H0 = numpyro.sample('H0', dist.Uniform(60, 80))
    Omega_m = numpyro.sample('Omega_m', dist.Uniform(0.1, 0.5))

    params = {'H0': H0, 'Omega_m': Omega_m}

    # Compute theoretical distance modulus
    grid = LCDM.compute_grid_traced(z_obs, params)
    d_L = grid['d_L']
    mu_theory = 5 * jnp.log10(d_L) + 25

    # Likelihood
    numpyro.sample('obs', dist.Normal(mu_theory, mu_err), obs=mu_obs)

Preset Parameters

HIcosmo provides commonly used cosmological parameter presets:

Planck 2018

Parameter	Value	Description
\(H_0\)	67.36	Hubble constant [km/s/Mpc]
\(\Omega_m\)	0.3153	Matter density parameter
\(\Omega_b\)	0.0493	Baryon density parameter
\(\sigma_8\)	0.8111	Matter power spectrum normalization
\(n_s\)	0.9649	Primordial power spectrum scalar spectral index

# Use Planck 2018 default parameters
model = LCDM()  # Default is Planck 2018

# Or explicitly use a preset
from hicosmo.parameters import ParameterRegistry
registry = ParameterRegistry.from_defaults('planck2018')

WMAP9

Parameter	Value	Description
\(H_0\)	70.0	Hubble constant [km/s/Mpc]
\(\Omega_m\)	0.279	Matter density parameter
\(\Omega_b\)	0.046	Baryon density parameter

Adding New Cosmological Models

HIcosmo adopts a minimalist model architecture design, allowing users to create new cosmological models by defining only the core physics formulas. This section describes in detail how to add custom models.

Design Principles

HIcosmo’s model architecture follows these core principles:

Principle	Description
Minimal subclass	Subclasses only define differential physics + necessary parameters
Unified signature	All core computation functions use the func(z, params) signature
User-first naming	Public API uses the intuitive name E_z
Decorator generation	Boilerplate code is auto-generated by decorators
Physics-framework separation	Subclasses handle physics only; framework logic lives in base class/decorator

Model Inheritance Hierarchy

CosmologyBase (abstract base class)
    |
    └── LCDM (reference implementation)
            |
            ├── wCDM (constant w)
            ├── CPL  (evolving w(z))
            └── ILCDM (interaction)

All custom models should inherit from LCDM or its subclasses.

Minimal Subclass Pattern

A minimal cosmological model requires only the following:

Must Define	May Define	Never Define
__init__	w_z() (if different)	_E_z_static
@staticmethod E_z(z, params)	get_parameters() (if new parameters exist)	Distance calculation methods
		Factory methods

Core idea: Only define the physics formulas that differ from the parent class; everything else is handled automatically by the decorator and base class.

Complete Example: Creating a Quintessence Model

Below we create a simple Quintessence dark energy model with the equation of state:

\[w(z) = w_0 + w_1 \ln(1+z)\]

Step 1: Create the model class

# hicosmo/models/quintessence.py

from typing import Dict, Union
import jax.numpy as jnp

from .lcdm import LCDM
from .base import register_cosmology_model


@register_cosmology_model
class Quintessence(LCDM):
    """
    Quintessence dark energy model.

    Equation of state parameterization: w(z) = w0 + w1 * ln(1+z)

    Physical background:
    - Effective equation of state under scalar field slow-roll approximation
    - w1 > 0: dark energy is "stiffer" in the early universe
    - w1 < 0: dark energy is "softer" in the early universe
    """

    def __init__(self, w0: float = -1.0, w1: float = 0.0, **kwargs):
        """
        Initialize the Quintessence model.

        Parameters
        ----------
        w0 : float
            Dark energy equation of state today (z=0), default -1.0
        w1 : float
            Logarithmic evolution coefficient of the equation of state, default 0.0
        **kwargs
            Other parameters passed to LCDM (H0, Omega_m, etc.)
        """
        super().__init__(w0=w0, w1=w1, **kwargs)

    @staticmethod
    def E_z(z: Union[float, jnp.ndarray], params: Dict) -> jnp.ndarray:
        """
        Compute the dimensionless Hubble parameter E(z) = H(z)/H0.

        This is the core of the model: just define the physics formula!

        Mathematical form:
        For w(z) = w0 + w1*ln(1+z), the dark energy density evolves as:
        f_DE(z) = (1+z)^{3(1+w0)} * exp(3*w1*[(1+z)*ln(1+z) - z])
        """
        z_arr = jnp.asarray(z)
        one_plus_z = 1.0 + z_arr

        # Get parameters from the params dictionary
        Omega_m = params['Omega_m']
        Omega_r = params.get('Omega_r', 0.0)
        Omega_k = params.get('Omega_k', 0.0)
        Omega_Lambda = params.get('Omega_Lambda', 1.0 - Omega_m - Omega_k - Omega_r)
        w0 = params.get('w0', -1.0)
        w1 = params.get('w1', 0.0)

        # Dark energy density evolution factor
        # Obtained by integrating w(z) = w0 + w1*ln(1+z)
        log_term = one_plus_z * jnp.log(one_plus_z) - z_arr
        f_DE = one_plus_z**(3.0 * (1.0 + w0)) * jnp.exp(3.0 * w1 * log_term)

        # E^2(z) = Omega_m*(1+z)^3 + Omega_r*(1+z)^4 + Omega_k*(1+z)^2 + Omega_Lambda*f_DE
        E_squared = (
            Omega_m * one_plus_z**3 +
            Omega_r * one_plus_z**4 +
            Omega_k * one_plus_z**2 +
            Omega_Lambda * f_DE
        )
        return jnp.sqrt(E_squared)

    def w_z(self, z: Union[float, jnp.ndarray]) -> Union[float, jnp.ndarray]:
        """
        Dark energy equation of state w(z) = w0 + w1 * ln(1+z).

        Only override this method when the equation of state differs from LCDM.
        """
        z_arr = jnp.asarray(z)
        w0 = self.params.get('w0', -1.0)
        w1 = self.params.get('w1', 0.0)
        return w0 + w1 * jnp.log(1.0 + z_arr)

    @classmethod
    def get_parameters(cls):
        """
        Return the model parameter list.

        Only override this method when the model introduces new parameters.
        """
        from ..parameters import Parameter

        # Get parent class parameters
        params = LCDM.get_parameters()

        # Add new parameters
        params.extend([
            Parameter(
                name='w0',
                value=-1.0,
                free=False,
                prior={'dist': 'uniform', 'min': -2.0, 'max': 0.0},
                latex_label=r'$w_0$',
                description='Dark energy equation of state at z=0'
            ),
            Parameter(
                name='w1',
                value=0.0,
                free=False,
                prior={'dist': 'uniform', 'min': -1.0, 'max': 1.0},
                latex_label=r'$w_1$',
                description='Logarithmic evolution coefficient'
            ),
        ])
        return params

Step 2: Register the model (optional)

If you want the model to be available in hicosmo.models, add it to hicosmo/models/__init__.py:

from .quintessence import Quintessence

__all__ = ['LCDM', 'wCDM', 'CPL', 'ILCDM', 'Quintessence']

Step 3: Use the new model

from hicosmo.models import Quintessence

# Create model instance
model = Quintessence(H0=67.36, Omega_m=0.3, w0=-0.9, w1=0.1)

# All distance methods are automatically available (inherited from LCDM)
d_L = model.luminosity_distance(1.0)
D_A = model.angular_diameter_distance(1.0)

# View the evolving equation of state
import jax.numpy as jnp
z = jnp.array([0.0, 0.5, 1.0, 2.0])
w = model.w_z(z)
print(f"w(z) = {w}")

# Use in MCMC
params = {'H0': 70.0, 'Omega_m': 0.3, 'w0': -0.9, 'w1': 0.1}
z_grid = jnp.linspace(0.01, 2.0, 100)
grid = Quintessence.compute_grid_traced(z_grid, params)
d_L_array = grid['d_L']

What the Decorator Does

The @register_cosmology_model decorator automatically generates the following for your model:

Generated Item	Description
_E_z_kernel	Converted from the user-defined E_z
_E_z_static	JIT-compiled version for high-performance computation
E_z(self, z)	Instance method wrapper, automatically passes self.params
compute_grid_traced	MCMC factory method, supports JAX tracer
sound_horizon_traced	Sound horizon computation factory method
sound_horizon_drag_traced	Drag epoch sound horizon factory method
recombination_redshift_traced	Recombination redshift computation factory method

Core advantage: You only write the physics formula; the decorator handles all framework code!

Parameter Dictionary Specification

The params in E_z(z, params) is a dictionary containing all cosmological parameters:

Required parameters:

params = {
    'Omega_m': 0.3,    # Required: matter density parameter
}

Optional parameters (with defaults):

params = {
    'Omega_m': 0.3,
    'Omega_r': 0.0,              # Radiation density, default 0
    'Omega_k': 0.0,              # Curvature, default 0 (flat)
    'Omega_Lambda': 0.7,         # Dark energy density, default 1-Omega_m-Omega_k-Omega_r
    # Model-specific parameters
    'w': -1.0,                   # wCDM
    'w0': -1.0, 'wa': 0.0,       # CPL
    'beta': 0.0,                 # ILCDM
}

Recommended pattern: Use params.get('key', default) to provide defaults:

Omega_r = params.get('Omega_r', 0.0)  # Use 0.0 if not present

Testing New Models

After creating a new model, write tests to verify its correctness:

# tests/test_quintessence.py

import pytest
import jax.numpy as jnp
from hicosmo.models import Quintessence, LCDM


class TestQuintessence:
    """Quintessence model test suite"""

    def test_reduces_to_lcdm(self):
        """Should reduce to LCDM when w0=-1, w1=0"""
        quint = Quintessence(H0=67.36, Omega_m=0.3, w0=-1.0, w1=0.0)
        lcdm = LCDM(H0=67.36, Omega_m=0.3)

        z = jnp.array([0.5, 1.0, 2.0])

        # E(z) should be identical
        E_quint = quint.E_z(z)
        E_lcdm = lcdm.E_z(z)
        assert jnp.allclose(E_quint, E_lcdm, rtol=1e-10)

        # Distances should also be identical
        d_L_quint = quint.luminosity_distance(1.0)
        d_L_lcdm = lcdm.luminosity_distance(1.0)
        assert jnp.isclose(d_L_quint, d_L_lcdm, rtol=1e-10)

    def test_equation_of_state(self):
        """Test equation of state w(z) = w0 + w1*ln(1+z)"""
        quint = Quintessence(w0=-0.9, w1=0.1)

        # At z=0, w = w0
        assert jnp.isclose(quint.w_z(0.0), -0.9, rtol=1e-10)

        # At z=e-1, w = w0 + w1*ln(e) = w0 + w1
        z_e = jnp.e - 1
        assert jnp.isclose(quint.w_z(z_e), -0.9 + 0.1, rtol=1e-6)

    def test_compute_grid_traced(self):
        """Test the traced interface"""
        params = {'H0': 70.0, 'Omega_m': 0.3, 'w0': -0.9, 'w1': 0.1}
        z_grid = jnp.linspace(0.01, 2.0, 100)

        grid = Quintessence.compute_grid_traced(z_grid, params)

        # Check return values
        assert 'd_L' in grid
        assert 'D_M' in grid
        assert 'E_z' in grid

        # Check shapes
        assert grid['d_L'].shape == z_grid.shape

    def test_jax_gradients(self):
        """Test JAX automatic differentiation"""
        from jax import grad

        def loss(H0):
            params = {'H0': H0, 'Omega_m': 0.3, 'w0': -1.0, 'w1': 0.0}
            grid = Quintessence.compute_grid_traced(jnp.array([1.0]), params)
            return grid['d_L'][0]

        # Gradient should exist and be finite
        grad_H0 = grad(loss)(70.0)
        assert jnp.isfinite(grad_H0)

Run tests:

pytest tests/test_quintessence.py -v

Common Errors and Solutions

Error 1: Forgetting the decorator

# Wrong: no decorator
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

# Correct: use the decorator
@register_cosmology_model
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

Error 2: ``E_z`` is not a static method

# Wrong: instance method
@register_cosmology_model
class MyModel(LCDM):
    def E_z(self, z, params):  # Should not have self
        ...

# Correct: static method
@register_cosmology_model
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

Error 3: Wrong parameter signature

# Wrong: expanded parameters
@staticmethod
def E_z(z, Omega_m, Omega_Lambda, w):
    ...

# Correct: use params dictionary
@staticmethod
def E_z(z, params):
    Omega_m = params['Omega_m']
    Omega_Lambda = params.get('Omega_Lambda', 1.0 - Omega_m)
    w = params.get('w', -1.0)
    ...

Error 4: Redefining framework methods in subclass

# Wrong: these methods do not need to be redefined
@register_cosmology_model
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

    @staticmethod
    def _E_z_static(z, params):  # The decorator generates this
        ...

    def compute_grid_traced(self, z, params):  # The decorator generates this
        ...

Performance Optimization

JIT Compilation

All models’ core computations are optimized with JAX JIT compilation:

from jax import jit
import jax.numpy as jnp
from hicosmo.models import LCDM

model = LCDM()

# First call: compilation
%timeit model.luminosity_distance(1.0)  # ~100ms (includes compilation)

# Subsequent calls: compiled
%timeit model.luminosity_distance(1.0)  # ~0.01ms

Vectorized Computation

Use vmap for efficient batch computation:

from jax import vmap

z_array = jnp.linspace(0.01, 2.0, 1000)

# Vectorized distance computation
d_L_array = vmap(model.luminosity_distance)(z_array)

Or use compute_grid_traced directly:

grid = LCDM.compute_grid_traced(z_array, {'H0': 70.0, 'Omega_m': 0.3})
d_L_array = grid['d_L']

Performance Benchmarks

Operation	scipy (qcosmc)	HIcosmo (JAX)	Speedup
Distance calculation (1000 points)	150 ms	20 ms	7.5x
E(z) computation (10000 points)	50 ms	2 ms	25x
Growth factor	120 ms	15 ms	8x

Best Practices

1. Choose the appropriate model

   • Use LCDM for standard analyses

   • Use CPL to explore dark energy properties

   • Use ILCDM to study dark matter–dark energy interaction

2. Use preset parameters

   # Use Planck priors
   model = LCDM()  # Default Planck 2018
   

3. Batch computation optimization

   # Recommended: compute multiple redshifts at once
   grid = LCDM.compute_grid_traced(z_array, params)
   
   # Avoid: loop calls
   for z in z_array:
       d_L = model.luminosity_distance(z)
   

4. Use traced interface for MCMC

   # Use static methods in NumPyro models
   grid = LCDM.compute_grid_traced(z_obs, params)
   

Next Steps

• Likelihood Functions: Learn how to use observational data

• Samplers: Learn about MCMC parameter estimation

• Fisher Forecasts: Perform survey predictions