配置驱动
HIcosmo 支持类似 Cobaya 的 YAML 配置。
并行初始化
在运行任何配置前,先初始化并行环境:
import hicosmo as hc
hc.init(8) # 8 个并行设备
备注
hc.init(N) 创建 N 个 JAX 逻辑设备,实现 N 条链真正并行。
hc.init(8)- 8 个并行设备(推荐)hc.init()- 自动检测(最多 8 个)hc.init("GPU")- GPU 模式
YAML 配置示例
name: joint_sn_bao
theory: LCDM
likelihood:
- name: sn
dataset: pantheon+
- name: bao
dataset: sdss_dr12
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: 1000
num_chains: 8 # 默认 = 设备数
num_warmup: 300
output:
root: results
chain_name: joint_sn_bao
运行配置
# 1. 先初始化
import hicosmo as hc
hc.init(8)
# 2. 运行配置
from hicosmo import run_from_config
run_from_config("examples/tutorials/configs/joint_sn_bao.yaml")