宇宙学模型

HIcosmo 提供多种宇宙学模型，涵盖标准宇宙学模型及其动态暗能量扩展。所有模型都基于 JAX 构建，支持 JIT 编译和自动微分。

模型概览

模型	暗能量状态方程	描述
LCDM	\(w = -1\)	标准 \(\Lambda\text{CDM}\) 模型，宇宙学常数
wCDM	\(w = w_0\)	常数暗能量状态方程
CPL	\(w(a) = w_0 + w_a(1-a)\)	Chevallier-Polarski-Linder 参数化
ILCDM	\(w = -1\) + 相互作用	相互作用暗能量模型

基础理论

无量纲 Hubble 参数

所有模型的核心是无量纲 Hubble 参数 \(E(z) = H(z)/H_0\)。它决定了宇宙的膨胀历史和所有距离测量。

通用形式：

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

其中：

• \(\Omega_m\)：物质密度参数

• \(\Omega_r\)：辐射密度参数

• \(\Omega_k\)：曲率密度参数

• \(\Omega_{\rm DE}(z)\)：暗能量密度演化

归一化条件：在 \(z=0\) 时，\(E(0) = 1\)，这要求：

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

距离计算

从 \(E(z)\) 可以推导所有宇宙学距离：

共动距离：

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

横向共动距离（考虑曲率）：

\[\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}\]

其中 \(D_H = c/H_0\) 是 Hubble 距离。

光度距离：

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

角直径距离：

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

距离模数：

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

LCDM 模型

标准 \(\Lambda\text{CDM}\) 模型假设暗能量是宇宙学常数，\(w = -1\)。

LCDM 数学形式

Hubble 参数：

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

暗能量状态方程：

\[w(z) = -1 \quad \text{（常数）}\]

LCDM 参数说明

参数	默认值	范围	描述
H0	67.36	[50, 100]	Hubble 常数，单位 km/s/Mpc
Omega_m	0.3153	[0.01, 0.99]	总物质密度参数
Omega_b	0.0493	[0.01, 0.1]	重子密度参数
Omega_k	0.0	[-0.3, 0.3]	曲率密度参数（0 表示平直宇宙）
Omega_r	自动计算	—	辐射密度参数，由 \(T_{\rm CMB}\) 计算
sigma8	0.8111	[0.5, 1.2]	物质功率谱归一化
n_s	0.9649	[0.8, 1.1]	原初功率谱标量指数
T_cmb	2.7255	—	CMB 温度 (K)
N_eff	3.046	—	有效中微子数目

LCDM 基本使用

from hicosmo.models import LCDM

# 使用 Planck 2018 默认参数
model = LCDM()

# 或指定参数
model = LCDM(H0=70.0, Omega_m=0.3, Omega_k=0.01)

# 计算距离
z = 1.0
d_L = model.luminosity_distance(z)      # 光度距离 [Mpc]
D_A = model.angular_diameter_distance(z)  # 角直径距离 [Mpc]
D_M = model.transverse_comoving_distance(z)  # 横向共动距离 [Mpc]
mu = model.distance_modulus(z)          # 距离模数

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}")

参数一致性与更新

所有模型的参数都以 model.params 为单一真相（内部为 CosmologicalParameters），模型属性 w / w0 / wa / beta 只是对 params 的轻量封装，保证实例方法与 compute_grid_traced 的物理计算一致。

推荐的参数更新方式：

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                 # 等价于 w_model.params['w'] = -0.9

cpl = CPL(w0=-1.0, wa=0.0)
cpl.w0 = -0.95                   # 等价于 cpl.params['w0'] = -0.95
cpl.wa = 0.2

ilcdm = ILCDM(beta=0.0)
ilcdm.beta = 0.05                # 等价于 ilcdm.params['beta'] = 0.05

声音视界计算

LCDM 模型包含拖拽时刻声音视界 \(r_d\) 的计算，这是 BAO 分析的关键量。

声音视界（Eisenstein & Hu 1998 拟合公式）：

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

# 获取声音视界
rd = model.sound_horizon()  # Mpc
rd_h = model.rd_h           # r_d × h (无量纲)

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

增长函数

LCDM 模型提供线性增长因子和增长率的计算。

增长因子 D(z)（Carroll, Press & Turner 1992 近似）：

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

增长率：

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

# 增长因子（归一化到 z=0）
D = model.growth_factor(z=0.5)

# 增长率
f = model.growth_rate(z=0.5)

# f*sigma8 参数（RSD 分析常用）
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 模型

wCDM 模型将暗能量状态方程推广为常数 \(w \neq -1\)。

wCDM 数学形式

Hubble 参数：

\[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)}\]

暗能量状态方程：

\[w(z) = w_0 \quad \text{（常数）}\]

wCDM 参数说明

wCDM 继承 LCDM 所有参数，并增加：

参数	默认值	先验范围	描述
w (或 w0)	-1.0	[-2.5, 0.0]	暗能量状态方程常数

物理约束：

• \(w = -1\)：宇宙学常数（LCDM）

• \(w > -1\)：Quintessence 类暗能量

• \(w < -1\)：Phantom 暗能量（违反能量条件）

wCDM 基本使用

from hicosmo.models import wCDM

# 创建 wCDM 模型
model = wCDM(H0=70.0, Omega_m=0.3, w=-0.9)

# 所有距离方法继承自 LCDM
d_L = model.luminosity_distance(1.0)

# 查看状态方程
w = model.w_z(0.5)  # 返回常数 w
print(f"w(z=0.5) = {w:.2f}")

# E(z) 现在包含 w 参数
E = model.E_z(1.0)
print(f"E(z=1) = {E:.4f}")

CPL 模型

CPL (Chevallier-Polarski-Linder) 参数化允许暗能量状态方程随时间演化。

CPL 数学形式

状态方程参数化：

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

其中 \(a = 1/(1+z)\) 是尺度因子。

暗能量密度演化：

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

Hubble 参数：

\[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 参数说明

CPL 继承 LCDM 所有参数，并增加：

参数	默认值	先验范围	描述
w0	-1.0	[-2.5, 0.5]	今日 (\(z=0\)) 的暗能量状态方程
wa	0.0	[-3.0, 2.0]	状态方程演化参数

特殊情况：

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

• \(w_a = 0\)：wCDM（常数 \(w\)）

• \(w_a \neq 0\)：动态暗能量

CPL 基本使用

from hicosmo.models import CPL

# 创建 CPL 模型
model = CPL(H0=67.36, Omega_m=0.3, w0=-1.0, wa=0.1)

# 查看演化的状态方程
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}")

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

# 距离计算
d_L = model.luminosity_distance(1.0)
print(f"d_L(z=1) = {d_L:.2f} Mpc")

暗能量品质因子

CPL 参数化常用于计算暗能量品质因子（Figure of Merit）：

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

FoM 越大，对暗能量状态方程的约束越好。

ILCDM 模型

ILCDM (Interacting Lambda-CDM) 模型假设暗物质和暗能量之间存在能量交换。

物理图像

在标准 LCDM 中，暗物质和暗能量各自独立演化。ILCDM 引入相互作用项：

\[Q = \beta H \rho_c\]

其中 \(\beta\) 是无量纲耦合常数，\(\rho_c\) 是冷暗物质密度。

能量守恒方程：

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

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

密度演化：

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

ILCDM 数学形式

Hubble 参数：

\[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]\]

其中 \(\Omega_c = \Omega_m - \Omega_b\) 是冷暗物质密度参数。

ILCDM 参数说明

参数	默认值	先验范围	描述
beta	0.0	[-1.0, 1.0]	暗能量-暗物质耦合常数

物理意义：

• \(\beta = 0\)：标准 LCDM

• \(\beta > 0\)：暗物质衰变为暗能量

• \(\beta < 0\)：暗能量衰变为暗物质

注意：ILCDM 假设平直宇宙（\(\Omega_k = 0\)）。

ILCDM 基本使用

from hicosmo.models import ILCDM

# 创建 ILCDM 模型
model = ILCDM(H0=67.36, Omega_m=0.3, beta=0.05)

# 距离计算
d_L = model.luminosity_distance(1.0)

# 比较不同 beta 值
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 接口

为了与 NumPyro MCMC 采样器兼容，所有模型都提供 compute_grid_traced 静态方法，支持 JAX tracer 传播。

基本用法

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

# 定义红移网格
z_grid = jnp.linspace(0.01, 2.0, 1000)

# 参数字典（用于 MCMC 采样）
# 只需提供 H0 / Omega_m 以及模型特有参数
params = {'H0': 70.0, 'Omega_m': 0.3}

# 批量计算所有距离
grid = LCDM.compute_grid_traced(z_grid, params)

# 返回字典包含多个量
d_L = grid['d_L']     # 光度距离
D_M = grid['D_M']     # 横向共动距离
D_H = grid['D_H']     # Hubble 距离 (c/H(z))
E_z = grid['E_z']     # 无量纲 Hubble 参数
d_C = grid['d_C']     # 共动距离
dVc_dz = grid['dVc_dz']  # 共动体积元

CPL 模型示例

# CPL 需要额外参数
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 参数字典示例

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)

在 NumPyro 中使用

import numpyro
import numpyro.distributions as dist

def model(z_obs, mu_obs, mu_err):
    # 采样宇宙学参数
    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}

    # 计算理论距离模数
    grid = LCDM.compute_grid_traced(z_obs, params)
    d_L = grid['d_L']
    mu_theory = 5 * jnp.log10(d_L) + 25

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

预设参数

HIcosmo 提供常用宇宙学参数预设：

Planck 2018

参数	值	说明
\(H_0\)	67.36	Hubble 常数 [km/s/Mpc]
\(\Omega_m\)	0.3153	物质密度参数
\(\Omega_b\)	0.0493	重子密度参数
\(\sigma_8\)	0.8111	物质功率谱归一化
\(n_s\)	0.9649	原初功率谱标量指数

# 使用 Planck 2018 默认参数
model = LCDM()  # 默认就是 Planck 2018

# 或显式使用预设
from hicosmo.parameters import ParameterRegistry
registry = ParameterRegistry.from_defaults('planck2018')

WMAP9

参数	值	说明
\(H_0\)	70.0	Hubble 常数 [km/s/Mpc]
\(\Omega_m\)	0.279	物质密度参数
\(\Omega_b\)	0.046	重子密度参数

添加新宇宙学模型

HIcosmo 采用极简的模型架构设计，让用户只需定义核心物理公式即可创建新的宇宙学模型。本节详细介绍如何添加自定义模型。

设计原则

HIcosmo 模型架构遵循以下核心原则：

原则	描述
最小化子类	子类只定义差异性物理 + 必要参数
统一签名	所有核心计算函数使用 func(z, params) 签名
用户优先命名	公开 API 使用直觉名称 E_z
装饰器生成	样板代码由装饰器自动生成
物理与框架分离	子类只管物理，框架逻辑在基类/装饰器

模型继承层次

CosmologyBase (抽象基类)
    │
    └── LCDM (参考实现)
            │
            ├── wCDM (常数 w)
            ├── CPL  (演化 w(z))
            └── ILCDM (相互作用)

所有自定义模型应继承自 LCDM 或其子类。

最小子类模式

一个最简的宇宙学模型只需要以下内容：

必须定义	可选定义	绝不定义
__init__	w_z() (如果不同)	_E_z_static
@staticmethod E_z(z, params)	get_parameters() (如果有新参数)	距离计算方法
		工厂方法

核心思想：只定义与父类不同的物理公式，其他一切由装饰器和基类自动处理。

完整示例：创建 Quintessence 模型

下面创建一个简单的 Quintessence 暗能量模型，其状态方程为：

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

步骤 1：创建模型类

# 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 暗能量模型。

    状态方程参数化：w(z) = w0 + w1 * ln(1+z)

    物理背景：
    - 标量场慢滚近似下的有效状态方程
    - w1 > 0：早期宇宙暗能量更"硬"
    - w1 < 0：早期宇宙暗能量更"软"
    """

    def __init__(self, w0: float = -1.0, w1: float = 0.0, **kwargs):
        """
        初始化 Quintessence 模型。

        Parameters
        ----------
        w0 : float
            今日 (z=0) 的暗能量状态方程，默认 -1.0
        w1 : float
            状态方程的对数演化系数，默认 0.0
        **kwargs
            传递给 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:
        """
        计算无量纲 Hubble 参数 E(z) = H(z)/H0。

        这是模型的核心：只需定义物理公式！

        数学形式：
        对于 w(z) = w0 + w1*ln(1+z)，暗能量密度演化为：
        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

        # 从 params 字典获取参数
        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)

        # 暗能量密度演化因子
        # 对 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]:
        """
        暗能量状态方程 w(z) = w0 + w1 * ln(1+z)。

        只有当状态方程与 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):
        """
        返回模型参数列表。

        只有当模型引入新参数时才需要重写此方法。
        """
        from ..parameters import Parameter

        # 获取父类参数
        params = LCDM.get_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

步骤 2：注册模型（可选）

如果想让模型在 hicosmo.models 中可用，需要在 hicosmo/models/__init__.py 中添加：

from .quintessence import Quintessence

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

步骤 3：使用新模型

from hicosmo.models import Quintessence

# 创建模型实例
model = Quintessence(H0=67.36, Omega_m=0.3, w0=-0.9, w1=0.1)

# 所有距离方法自动可用（继承自 LCDM）
d_L = model.luminosity_distance(1.0)
D_A = model.angular_diameter_distance(1.0)

# 查看演化的状态方程
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}")

# 在 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']

装饰器的作用

@register_cosmology_model 装饰器会自动为你的模型生成：

生成的内容	说明
_E_z_kernel	从用户定义的 E_z 转换而来
_E_z_static	JIT 编译版本，用于高性能计算
E_z(self, z)	实例方法包装，自动传入 self.params
compute_grid_traced	MCMC 工厂方法，支持 JAX tracer
sound_horizon_traced	声音视界计算工厂方法
sound_horizon_drag_traced	拖拽时刻声视界工厂方法
recombination_redshift_traced	重组红移计算工厂方法

核心优势：你只需写物理公式，装饰器处理所有框架代码！

参数字典规范

E_z(z, params) 中的 params 是一个字典，包含所有宇宙学参数：

必须参数：

params = {
    'Omega_m': 0.3,    # 必须：物质密度参数
}

**可选参数**（有默认值）：

params = {
    'Omega_m': 0.3,
    'Omega_r': 0.0,              # 辐射密度，默认 0
    'Omega_k': 0.0,              # 曲率，默认 0（平直）
    'Omega_Lambda': 0.7,         # 暗能量密度，默认 1-Omega_m-Omega_k-Omega_r
    # 模型特有参数
    'w': -1.0,                   # wCDM
    'w0': -1.0, 'wa': 0.0,       # CPL
    'beta': 0.0,                 # ILCDM
}

推荐写法：使用 params.get('key', default) 提供默认值：

Omega_r = params.get('Omega_r', 0.0)  # 如果不存在，使用 0.0

测试新模型

创建新模型后，应编写测试验证其正确性：

# tests/test_quintessence.py

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


class TestQuintessence:
    """Quintessence 模型测试套件"""

    def test_reduces_to_lcdm(self):
        """w0=-1, w1=0 时应退化为 LCDM"""
        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) 应该相同
        E_quint = quint.E_z(z)
        E_lcdm = lcdm.E_z(z)
        assert jnp.allclose(E_quint, E_lcdm, rtol=1e-10)

        # 距离也应该相同
        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):
        """测试状态方程 w(z) = w0 + w1*ln(1+z)"""
        quint = Quintessence(w0=-0.9, w1=0.1)

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

        # 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):
        """测试 traced 接口"""
        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)

        # 检查返回值
        assert 'd_L' in grid
        assert 'D_M' in grid
        assert 'E_z' in grid

        # 检查形状
        assert grid['d_L'].shape == z_grid.shape

    def test_jax_gradients(self):
        """测试 JAX 自动微分"""
        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]

        # 梯度应该存在且有限
        grad_H0 = grad(loss)(70.0)
        assert jnp.isfinite(grad_H0)

运行测试：

pytest tests/test_quintessence.py -v

常见错误与解决

错误 1：忘记使用装饰器

# ❌ 错误：没有装饰器
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

# ✅ 正确：使用装饰器
@register_cosmology_model
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

错误 2：``E_z`` 不是静态方法

# ❌ 错误：实例方法
@register_cosmology_model
class MyModel(LCDM):
    def E_z(self, z, params):  # 不应该有 self
        ...

# ✅ 正确：静态方法
@register_cosmology_model
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

错误 3：参数签名错误

# ❌ 错误：展开参数
@staticmethod
def E_z(z, Omega_m, Omega_Lambda, w):
    ...

# ✅ 正确：使用 params 字典
@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)
    ...

错误 4：在子类中重复定义框架方法

# ❌ 错误：不需要重新定义这些方法
@register_cosmology_model
class MyModel(LCDM):
    @staticmethod
    def E_z(z, params):
        ...

    @staticmethod
    def _E_z_static(z, params):  # ❌ 装饰器会生成
        ...

    def compute_grid_traced(self, z, params):  # ❌ 装饰器会生成
        ...

性能优化

JIT 编译

所有模型的核心计算都经过 JAX JIT 编译优化：

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

model = LCDM()

# 首次调用：编译
%timeit model.luminosity_distance(1.0)  # ~100ms (含编译)

# 后续调用：已编译
%timeit model.luminosity_distance(1.0)  # ~0.01ms

向量化计算

使用 vmap 实现高效批量计算：

from jax import vmap

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

# 向量化距离计算
d_L_array = vmap(model.luminosity_distance)(z_array)

或者直接使用 compute_grid_traced：

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

性能基准

操作	scipy (qcosmc)	HIcosmo (JAX)	加速比
距离计算 (1000点)	150 ms	20 ms	7.5x
E(z) 计算 (10000点)	50 ms	2 ms	25x
增长因子	120 ms	15 ms	8x

最佳实践

1. 选择合适的模型

   • 标准分析使用 LCDM

   • 探索暗能量性质使用 CPL

   • 研究暗物质-暗能量相互作用使用 ILCDM

2. 利用预设参数

   # 使用 Planck 先验
   model = LCDM()  # 默认 Planck 2018
   

3. 批量计算优化

   # ✅ 推荐：一次计算多个红移
   grid = LCDM.compute_grid_traced(z_array, params)
   
   # ❌ 避免：循环调用
   for z in z_array:
       d_L = model.luminosity_distance(z)
   

4. MCMC 使用 traced 接口

   # 在 NumPyro 模型中使用静态方法
   grid = LCDM.compute_grid_traced(z_obs, params)
   

下一步

• 似然函数：了解如何使用观测数据

• 采样器：学习 MCMC 参数估计

• Fisher 预报：进行巡天预测