高斯隐马尔可夫模型（HMM）：前向–后向与 EM

# 每个隐状态 k 有自己的协方差矩阵 Σ_k（均值设为 0，即“连接状态”），状态序列服从一阶马尔可夫链。
# 本例用 numpy 写出 Baum–Welch（EM）算法：E 步用缩放的前向–后向递推求后验 γ_t(k)，
# M 步更新初始概率、转移矩阵和各状态协方差。模拟数据按 S1(300)→S2(200)→S1(300)→S2(200) 切换。
# 输出解读：后验最大状态与真值的一致率；估计的转移矩阵对角线接近 1，对应较长的停留时间。
# 真实研究用 HMM-MAR、osl-dynamics 等工具，并需要多次随机初始化和状态数敏感性分析。
import numpy as np
from scipy.stats import multivariate_normal

rng = np.random.default_rng(11)
N = 5
S1 = np.eye(N); S1[0, 1] = S1[1, 0] = S1[2, 3] = S1[3, 2] = 0.8
S2 = np.eye(N); S2[0, 4] = S2[4, 0] = S2[1, 2] = S2[2, 1] = -0.7
segs = [(0, 300), (1, 200), (0, 300), (1, 200)]
X = np.vstack([rng.multivariate_normal(np.zeros(N), [S1, S2][s], size=L) for s, L in segs])
truth = np.concatenate([[s] * L for s, L in segs])
T, K = X.shape[0], 2

def fit_hmm(X, K, n_iter=50, seed=0):
    r = np.random.default_rng(seed)
    T, N = X.shape
    pi = np.full(K, 1 / K)
    A = np.full((K, K), 0.05 / (K - 1)); np.fill_diagonal(A, 0.95)
    blocks = np.array_split(r.permutation(T), K)        # 随机划分时间点作初始协方差
    Sig = [np.cov(X[b].T) + 1e-3 * np.eye(N) for b in blocks]
    for _ in range(n_iter):
        B = np.column_stack([multivariate_normal(np.zeros(N), S).pdf(X) for S in Sig])  # (T, K)
        alpha = np.zeros((T, K)); c = np.zeros(T)
        alpha[0] = pi * B[0]; c[0] = alpha[0].sum(); alpha[0] /= c[0]
        for t in range(1, T):
            alpha[t] = (alpha[t - 1] @ A) * B[t]; c[t] = alpha[t].sum(); alpha[t] /= c[t]
        beta = np.ones((T, K))
        for t in range(T - 2, -1, -1):
            beta[t] = (A @ (B[t + 1] * beta[t + 1])) / c[t + 1]
        gamma = alpha * beta
        xi = np.zeros((K, K))
        for t in range(T - 1):
            xi += (alpha[t][:, None] * A * (B[t + 1] * beta[t + 1])[None, :]) / c[t + 1]
        pi = gamma[0]
        A = xi / xi.sum(1, keepdims=True)
        Sig = [(gamma[:, k, None] * X).T @ X / gamma[:, k].sum() + 1e-6 * np.eye(N) for k in range(K)]
    return gamma, A, Sig, np.log(c).sum()

gamma, A, Sig, ll = fit_hmm(X, K)
path = gamma.argmax(1)
acc = max(np.mean(path == truth), np.mean(path != truth))
print(f"对数似然 = {ll:.1f}；后验最大状态与真值一致率 = {acc:.3f}")
print("估计的转移矩阵：\n", A.round(3))
print("平均停留时间（时间点）≈ 1/(1 − A_kk)：", (1 / (1 - np.diag(A))).round(1))

import matplotlib.pyplot as plt
k1 = int(np.argmax([np.mean(path[truth == 0] == k) for k in range(K)]))   # 与 S1 对应的估计状态
fig, axes = plt.subplots(1, 3, figsize=(11, 3.6), gridspec_kw={"width_ratios": [2.4, 1, 1]})
axes[0].fill_between(np.arange(T), 0, 1, where=truth == 0, color="#DDDDDD", step="mid", label="true S1 periods")
axes[0].plot(gamma[:, k1], color="#4477AA", lw=0.8, label="P(S1 | data)")
axes[0].set(xlabel="Time point", ylabel="Posterior probability", ylim=(-0.02, 1.02),
            title=f"HMM state posterior (agreement {acc:.2f})")
axes[0].legend(fontsize=8, loc="lower right")
for ax, k, nm in [(axes[1], k1, "S1"), (axes[2], 1 - k1, "S2")]:
    sd = np.sqrt(np.diag(Sig[k]))
    im = ax.imshow(Sig[k] / np.outer(sd, sd), cmap="RdBu_r", vmin=-1, vmax=1)
    ax.set(title=f"{nm} (estimated r)", xticks=range(N), yticks=range(N))
fig.colorbar(im, ax=axes[2], shrink=0.8)
plt.tight_layout(); plt.show()