Fix RegimeSwitchingModel: MEADS+GHMC with k-fold adaptation (Rhat ≤ 1.01)

[junpenglao]

Summary

• Numerically stable HMM forward filter: replaced linear-space norm.pdf with a fully log-space implementation (norm.logpdf, logaddexp, logsumexp) to avoid underflow in the filtered state probabilities.
• Fixed MEADS divergence via two-stage initialization: the original code initialised all chains at nearly the same point, making inter-chain gradient covariance ≈ 0, causing MEADS to set a huge step_size and GHMC chains to diverge (Rhat ~ 1e14). Fix:
  1. Anchor sigma/rho to break regime symmetry (sigma=[3,10] in constrained space).
  2. Run a single NUTS chain (2000 warm-up steps) to reach the posterior.
  3. Spawn 64 short NUTS chains (200 steps each) from the warm endpoint to generate diverse starting positions that span the true posterior width — giving MEADS a well-conditioned gradient covariance estimate.
• Sampler: MEADS k-fold (num_folds=4, Algorithm 3 of Hoffman & Sountsov 2022) + GHMC, 64 chains × 1000 draws.

Results

Parameter	Rhat	n_eff
alpha[0,1]	1.01 / 1.00	~10k
sigma[0,1]	1.01 / 1.01	~9k / ~6k
p[0,1]	1.01 / 1.01	~7k / ~9k
rho	1.00	~10k
xi_0, y_0	1.00 / 1.01	~9k / ~7k

Runtime: ~27s (2× A100 or equivalent).

Test plan

• Run notebook end-to-end and verify Rhat ≤ 1.05 for all parameters
• Verify pair plot shows well-separated posteriors for the two regimes

🤖 Generated with Claude Code