如何通过 ceffyl.Sampler 的 PTMCMC 结果计算 Bayes evidence（logZ）

# run_pt.py
from ceffyl.Sampler import Sampler

# 需要保证 pta_NG15 已在脚本里正确构造
# 例如：pta_NG15 = ...

outdir = "./mcmc2"

sampler = Sampler.setup_sampler(
    pta_NG15,
    outdir=outdir,
    logL=pta_NG15.ln_likelihood,   # 或者 pta_NG15.get_lnlikelihood 的封装
    logp=pta_NG15.ln_prior,        # MCMC 场景使用 ln_prior（不是 prior transform）
    resume=False,
    jump=True,
)

x0 = pta_NG15.initial_samples()

sampler.sample(
    x0,
    int(20_000_000),
    Tmin=1,
    Tmax=1e3,               # 可调：见 2.3 节
    Tskip=100,
    isave=1000,
    thin=10,
    writeHotChains=True,    # 关键：写出热链
)

import la_forge.core as co

outdir = "./mcmc2"
c = co.Core(outdir, burn=0.85, pt_chains=True)

# 参数列名列表
print("params:", c.params)
print("has hot_chains:", bool(c.hot_chains))
if c.hot_chains:
    print("hot chain keys:", list(c.hot_chains.keys())[:10])

import numpy as np
import la_forge.core as co

outdir = "./mcmc2"
c = co.Core(outdir, burn=0.85, pt_chains=True)
lnlike_idx = c.params.index("lnlike")

pairs = []

# 冷链：T=1 -> beta=1
pairs.append((1.0, c.chain[c.burn:, lnlike_idx]))

# 热链：key 常是温度或 beta 的字符串（视版本而定）
for key, ch in (c.hot_chains or {}).items():
    # 常见情况：key 是温度 T
    T = float(key)
    beta = 1.0 / T
    pairs.append((beta, ch[c.burn:, lnlike_idx]))

# beta 升序
pairs.sort(key=lambda t: t[0])
betas = np.array([b for b, _ in pairs])
lnL_series = [x for _, x in pairs]

print("beta_min =", betas[0], "T_max ~", 1.0/betas[0])
print("N_temps =", len(betas))

import numpy as np
from scipy.integrate import simpson

# 计算每个 beta 的均值
ElogL = np.array([x.mean() for x in lnL_series])

# （推荐）只对已覆盖的 beta 做积分；不要在缺少高温端时“硬补 beta=0”
logZ_ti = simpson(ElogL, x=betas)
print("logZ (TI, covered beta-range) =", logZ_ti)

import numpy as np

# 假设存在 prior_transform(u) -> theta
# u: shape (d,) in [0,1]
# theta: parameter vector compatible with pta.get_lnlikelihood

def estimate_ElogL_beta0(pta, prior_transform, nsamp=20000, seed=0):
    rng = np.random.default_rng(seed)
    d = len(pta.param_names)
    vals = []
    for _ in range(nsamp):
        u = rng.random(d)
        theta = prior_transform(u)
        vals.append(pta.get_lnlikelihood(theta))
    return float(np.mean(vals)), float(np.std(vals, ddof=1)/np.sqrt(nsamp))

E0, E0_se = estimate_ElogL_beta0(pta_NG15, prior_transform=pta_NG15.prior_transform)
print("E[lnL] at beta=0 ~", E0, "+/-", E0_se)

from scipy.integrate import simpson
import numpy as np

ElogL = np.array([x.mean() for x in lnL_series])
if betas[0] > 0:
    betas2 = np.r_[0.0, betas]
    ElogL2 = np.r_[E0, ElogL]   # 用 prior 采样给出的 E0
else:
    betas2, ElogL2 = betas, ElogL

logZ_ti_full = simpson(ElogL2, x=betas2)
print("logZ (TI, with beta=0 prior estimate) =", logZ_ti_full)

import numpy as np
from scipy.special import logsumexp

def logZ_stepping_stone(betas, lnL_series, logZ_beta0=0.0):
    # betas: ascending array, should include 1.0; can optionally include 0.0
    # lnL_series[k]: samples of lnL at beta=betas[k]
    betas = np.asarray(betas)
    logZ = float(logZ_beta0)

    for k in range(1, len(betas)):
        db = betas[k-1] - betas[k]  # negative number if betas ascending
        lnL = np.asarray(lnL_series[k])
        # E_{beta_k}[exp(db * lnL)]  (db <= 0)
        log_mean = logsumexp(db * lnL) - np.log(len(lnL))
        logZ += -log_mean

    return float(logZ)

# 如果 betas 不含 0，可把 beta=0 的 prior 采样加入（推荐）
# 例如：betas2 = [0] + betas, lnL_series2 = [prior_lnL_samples] + lnL_series

logZ_ss = logZ_stepping_stone(betas, lnL_series, logZ_beta0=0.0)
print("logZ (SS, covered beta-range) =", logZ_ss)

import numpy as np
import la_forge.core as co

outdir = "./mcmc2"
c = co.Core(outdir, burn=0.85, pt_chains=True)

lnlike_idx = c.params.index("lnlike")

# 建议使用 pta.param_names 的顺序与链对齐
pta_names = getattr(pta_NG15, "param_names", None)
if pta_names is None:
    pta_names = [par.name for par in pta_NG15.params]

col_idx = [c.params.index(n) for n in pta_names]

rng = np.random.default_rng(0)
rows = rng.integers(low=c.burn, high=c.chain.shape[0], size=20)

diffs = []
for r in rows:
    theta = c.chain[r, col_idx]
    ll_chain = c.chain[r, lnlike_idx]
    ll_eval  = pta_NG15.get_lnlikelihood(theta)
    diffs.append(ll_chain - ll_eval)

diffs = np.array(diffs)
print("lnlike(chain)-lnlike(pta) mean =", diffs.mean(), "std =", diffs.std())

import numpy as np
from scipy.integrate import simpson

def block_resample(x, B, rng):
    x = np.asarray(x)
    n = len(x)
    m = n // B
    x = x[:m*B].reshape(B, m)
    idx = rng.integers(0, B, size=B)
    return x[idx, :].reshape(-1)

def bootstrap_ti(betas, lnL_series, B=50, R=300, seed=0):
    rng = np.random.default_rng(seed)
    betas = np.asarray(betas)
    vals = []
    for _ in range(R):
        E = []
        for x in lnL_series:
            xr = block_resample(x, B, rng)
            E.append(xr.mean())
        vals.append(simpson(np.array(E), x=betas))
    vals = np.array(vals)
    return float(vals.mean()), float(vals.std(ddof=1))

m, s = bootstrap_ti(betas, lnL_series, B=50, R=300)
print("TI logZ ~", m, "+/-", s)