快速开始
本指南将带你快速上手 HIcosmo,实现宇宙学参数估计。
3 行代码完成推断
HIcosmo 的设计哲学是 3 行代码完成核心功能:
from hicosmo import hicosmo
inf = hicosmo(cosmology="LCDM", likelihood="sn", free_params=["H0", "Omega_m"])
samples = inf.run()
运行完毕后,你将获得 \(H_0\) 和 \(\Omega_m\) 的后验分布样本。
小技巧
并行加速:在脚本最开头添加以下代码启用多设备并行:
import hicosmo as hc
hc.init(8) # 8 个并行设备 = 8 条并行链 = ~800% CPU
必须在导入其他模块之前调用!
基本示例
单探针分析(超新星)
from hicosmo import hicosmo
# 使用 Pantheon+ 超新星数据约束 ΛCDM 模型
inf = hicosmo(
cosmology="LCDM",
likelihood="sn", # Pantheon+ 数据
free_params=["H0", "Omega_m"],
num_samples=4000, # 总样本数
num_chains=4 # 并行链数
)
# 运行 MCMC 采样
samples = inf.run()
# 查看结果摘要
inf.summary()
# 生成 corner plot
inf.corner_plot('corner_sn.pdf')
多探针联合分析
from hicosmo import hicosmo
# 联合 SN + BAO + CMB 约束
inf = hicosmo(
cosmology="LCDM",
likelihood=["sn", "bao", "cmb"], # 多探针联合
free_params=["H0", "Omega_m"],
num_samples=8000,
num_chains=4
)
samples = inf.run()
inf.summary()
inf.corner_plot('corner_joint.pdf')
支持的 Likelihood 字符串
HIcosmo 支持以下便捷的 likelihood 字符串:
字符串 |
对应类 |
说明 |
|---|---|---|
|
|
Pantheon+ 1701个 SNe Ia |
|
|
含 Cepheid 标校 |
|
|
DESI 2024 BAO |
|
|
SDSS DR12 BAO |
|
|
Planck 2018 压缩似然 |
|
|
强透镜时延 |
|
|
TDCOSMO 层次贝叶斯 |
|
|
SH0ES 局部 \(H_0\) |
扩展暗能量模型
**wCDM 模型**(常数暗能量状态方程):
inf = hicosmo(
cosmology="wCDM",
likelihood=["sn", "bao", "cmb"],
free_params=["H0", "Omega_m", "w0"],
num_samples=10000
)
samples = inf.run()
**CPL 参数化**(时变暗能量):
inf = hicosmo(
cosmology="CPL",
likelihood=["sn", "bao", "cmb"],
free_params=["H0", "Omega_m", "w0", "wa"],
num_samples=16000
)
samples = inf.run()
自定义似然函数
HIcosmo 不仅可以用于宇宙学分析,还可以用于任意参数拟合问题。以下示例展示如何拟合二次多项式 \(y = ax^2 + bx + c\):
#!/usr/bin/env python
"""
Simple MCMC Example: Polynomial Fitting (y = a*x² + b*x + c)
"""
import hicosmo as hc
hc.init()
import numpy as np
import jax.numpy as jnp
from pathlib import Path
from hicosmo.samplers import MCMC
from hicosmo.visualization import Plotter
# Load data
data_path = Path(__file__).parent / 'data' / 'sim_data.txt'
x, y_obs, y_err = np.loadtxt(data_path, unpack=True)
x, y_obs, y_err = jnp.asarray(x), jnp.asarray(y_obs), jnp.asarray(y_err)
def log_likelihood(a, b, c):
"""Log-likelihood: -0.5 * chi²"""
y_th = a * x**2 + b * x + c
return -0.5 * jnp.sum((y_obs - y_th)**2 / y_err**2)
# Parameters: {name: (initial, min, max, latex)}
params = {
'a': (3.5, 0.0, 10.0, r'$a$'),
'b': (1.0, 0.0, 4.0, r'$b$'),
'c': (1.0, 0.0, 3.0, r'$c$'),
}
if __name__ == '__main__':
# Run MCMC
mcmc = MCMC(params, log_likelihood, chain_name='polynomial_fit')
mcmc.run(num_samples=20000)
# Plot results
plotter = Plotter('polynomial_fit')
plotter.corner()
plotter.report()
参数格式说明:
# 字典格式: {name: (initial, min, max, latex)}
params = {
'a': (3.5, 0.0, 10.0, r'$a$'), # 参数 a: 初值 3.5,范围 [0, 10]
'b': (1.0, 0.0, 4.0, r'$b$'), # 参数 b: 初值 1.0,范围 [0, 4]
'c': (1.0, 0.0, 3.0, r'$c$'), # 参数 c: 初值 1.0,范围 [0, 3]
}
运行结果:
INFO:hicosmo.samplers.init:✅ HIcosmo: 8 devices ready for parallel MCMC
INFO:hicosmo.samplers.inference:MCMC completed:
INFO:hicosmo.samplers.inference: Start time: 2026-01-08 16:54:43
INFO:hicosmo.samplers.inference: End time: 2026-01-08 16:54:45
INFO:hicosmo.samplers.inference: Duration: 1.94s
INFO:hicosmo.samplers.inference: Saved to: mcmc_chains/polynomial_fit.h5
小技巧
生成的报告 (plotter.report()) 包含:
参数约束(均值、中位数、68%/95% 置信区间)
LaTeX 格式输出(可直接用于论文)
MCMC 诊断(ESS、R̂)
模型选择准则(AIC、贝叶斯证据)
使用类接口(高级)
对于需要更多控制的场景,可以直接使用类接口:
# 1. 首先初始化(必须在最前面!)
import hicosmo as hc
hc.init(8) # 8 个并行设备
# 2. 导入其他模块
from hicosmo.models import LCDM
from hicosmo.likelihoods import SN_likelihood, BAO_likelihood
from hicosmo.samplers import MCMC
# 3. 创建 likelihood 对象
sne = SN_likelihood(LCDM, "pantheon+", M_B="marginalize")
bao = BAO_likelihood(LCDM, "desi2024")
# 联合 likelihood
joint_likelihood = sne + bao
# 4. 配置参数
params = {
"H0": (70.0, 60.0, 80.0), # (参考值, 最小值, 最大值)
"Omega_m": (0.3, 0.1, 0.5),
}
# 5. 创建 MCMC 并运行(num_chains 默认 = 设备数 = 8)
mcmc = MCMC(params, joint_likelihood, chain_name="lcdm_sn_bao")
samples = mcmc.run(num_samples=4000)
# 6. 输出结果
mcmc.print_summary()
可视化结果
Corner Plot:
from hicosmo.visualization import Plotter
# 从保存的链加载
plotter = Plotter("lcdm_sn_bao")
plotter.corner(["H0", "Omega_m"], filename="corner.pdf")
多链对比:
plotter = Plotter(
["chain_planck", "chain_shoes", "chain_bao"],
chain_labels=["Planck", "SH0ES", "DESI BAO"]
)
plotter.corner(["H0", "Omega_m"], filename="h0_tension.pdf")
保存和加载结果
# 保存结果
inf.save_results("results/my_analysis.pkl")
# 加载并继续分析
from hicosmo.visualization import Plotter
plotter = Plotter("results/my_analysis.pkl")
plotter.corner(["H0", "Omega_m"])
YAML 配置驱动
对于可复现的分析,推荐使用 YAML 配置文件:
config.yaml:
name: joint_analysis
theory: LCDM
likelihood:
- name: sn
dataset: pantheon+
- name: bao
dataset: desi2024
- name: cmb
dataset: planck2018
params:
H0:
prior: {dist: uniform, min: 50, max: 100}
ref: 70.0
Omega_m:
prior: {dist: uniform, min: 0.1, max: 0.5}
ref: 0.3
sampler:
name: numpyro
num_samples: 8000
num_chains: 4
num_warmup: 1000
output:
root: results
chain_name: joint_lcdm
运行配置:
from hicosmo import run_from_config
run_from_config("config.yaml")