BenchMCMCDiag.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -fno-full-laziness -fno-cse #-}
-- | Diagnostic runner for the B7 MCMC bench (Phase B10b investigation).
--
-- Runs hanalyze NUTS on the 8-schools model with progressively
-- different configurations to localise the cause of poor ESS:
--
--   * default          : nutsStepSize=0.08, adapt=on (current)
--   * smaller-step     : nutsStepSize=0.02, adapt=off
--   * longer-warmup    : burnin 2000 (more dual-averaging)
--   * deeper-tree      : nutsMaxDepth=12 (default 10)
--
-- For each, prints accept rate, ESS(mu), ESS(tau), n samples > 1
-- distinct.
module Main where

import qualified Data.Map.Strict      as Map
import qualified Data.Text            as T
import qualified System.Random.MWC    as MWC
import qualified Data.Time.Clock      as Time
import           System.IO            (hSetBuffering, stdout, BufferMode (..))
import           Text.Printf          (printf)

import           Hanalyze.Model.HBM            (Distribution (..), ModelP, sample,
                                       observe)
import           Hanalyze.MCMC.Core            (Chain, chainAccepted, chainTotal,
                                       chainVals, posteriorMean,
                                       posteriorSD)
import           Hanalyze.MCMC.NUTS            (NUTSConfig (..), defaultNUTSConfig,
                                       nuts)
import           Hanalyze.Stat.MCMC            (ess)

schoolData :: [[Double]]
schoolData =
  [ [72, 68, 75, 71]
  , [85, 88, 82, 90]
  , [61, 65, 58, 63]
  ]

sigmaY :: Double
sigmaY = 5.0

schoolModel :: ModelP ()
schoolModel = do
  mu  <- sample "mu"  (Normal 0 100)
  tau <- sample "tau" (Exponential 0.1)
  mapM_ (\(j, ys) -> do
    theta <- sample (T.pack ("theta_" ++ show (j :: Int)))
                    (Normal mu tau)
    observe (T.pack ("y_" ++ show j))
            (Normal theta (realToFrac sigmaY)) ys)
    (zip [1 ..] schoolData)

initParams :: Map.Map T.Text Double
initParams = Map.fromList
  [ ("mu",      73.0)
  , ("tau",     10.0)
  , ("theta_1", 71.5)
  , ("theta_2", 86.25)
  , ("theta_3", 61.75)
  ]

runOne :: String -> NUTSConfig -> IO ()
runOne label cfg = do
  g <- MWC.create
  t0 <- Time.getCurrentTime
  ch <- nuts schoolModel cfg initParams g
  t1 <- Time.getCurrentTime
  let dt = realToFrac (Time.diffUTCTime t1 t0) :: Double
  reportChain label dt ch

reportChain :: String -> Double -> Chain -> IO ()
reportChain label dt ch = do
  let muV    = chainVals "mu"  ch
      tauV   = chainVals "tau" ch
      muEss  = ess muV
      tauEss = ess tauV
      acc    = fromIntegral (chainAccepted ch)
             / max 1 (fromIntegral (chainTotal ch)) :: Double
      muMean = maybe 0 id (posteriorMean "mu" ch)
      muSD   = maybe 0 id (posteriorSD   "mu" ch)
      tauMean = maybe 0 id (posteriorMean "tau" ch)
      muDistinct  = length (uniq muV)
      tauDistinct = length (uniq tauV)
      uniq xs = go xs []
        where go [] acc' = reverse acc'
              go (x:xs') acc'
                | x `elem` acc' = go xs' acc'
                | otherwise     = go xs' (x:acc')
  printf "=== %s ===\n" label
  printf "  time         %.2f s\n" dt
  printf "  accept       %.3f\n" acc
  printf "  mu  mean=%.3f sd=%.3f ess=%.1f distinct=%d\n"
    muMean muSD muEss muDistinct
  printf "  tau mean=%.3f                ess=%.1f distinct=%d\n"
    tauMean tauEss tauDistinct
  putStrLn ""

main :: IO ()
main = do
  hSetBuffering stdout LineBuffering
  putStrLn "================================================="
  putStrLn "  NUTS on 8-schools — diagnostic (B10b)"
  putStrLn "  Goal: explain ESS(mu)=42 vs blackjax ess=810"
  putStrLn "================================================="
  putStrLn ""

  let baseCfg = defaultNUTSConfig
        { nutsIterations = 1000
        , nutsBurnIn     = 500
        , nutsStepSize   = 0.08
        , nutsMaxDepth   = 10
        , nutsAdaptStepSize = True
        , nutsTargetAccept  = 0.8
        }

  -- Reduced iterations for fast diagnostic. ESS scales linearly with
  -- iterations so ratios stay informative.
  let cfg = baseCfg { nutsIterations = 200, nutsBurnIn = 100 }

  -- 1. baseline
  runOne "baseline (eps=0.08, adapt=on)" cfg

  -- 2. small step, no adapt
  runOne "small-step (eps=0.02, adapt=off)"
    cfg { nutsStepSize = 0.02, nutsAdaptStepSize = False }

  -- 3. high target accept (forces smaller eps)
  runOne "high-target (eps=0.08, target=0.95)"
    cfg { nutsTargetAccept = 0.95 }

  -- 4. shallower tree (limit search space)
  runOne "shallow-tree (maxDepth=5)"
    cfg { nutsMaxDepth = 5 }

  -- 5. full-size 1000 samples WITH diagonal mass-matrix adaptation (B11)
  runOne "full-size (1000 samples, mass adapt ON)" baseCfg
    { nutsIterations = 1000, nutsBurnIn = 500, nutsAdaptMass = True }

  -- 6. full-size with mass adapt OFF (baseline for comparison)
  runOne "full-size (1000 samples, mass adapt OFF)" baseCfg
    { nutsIterations = 1000, nutsBurnIn = 500, nutsAdaptMass = False }

  -- 7. mass adapt ON, longer warmup (2000) — does multi-window
  -- adaptation eventually converge given enough budget?
  runOne "long-warmup (1000 samples, warmup=2000, mass ON)" baseCfg
    { nutsIterations = 1000, nutsBurnIn = 2000, nutsAdaptMass = True }