imp-ppl-0.1.0.0: bench/Bench.hs

{-# OPTIONS_GHC -fno-full-laziness #-}
{-# LANGUAGE GADTs #-}
module Bench
  ( BenchCase(..)
  , BenchConfig(..)
  , defaultConfig
  , runAll
  ) where

import Control.Monad (replicateM, when)
import Data.List (sort)
import Data.Time.Clock (getCurrentTime, diffUTCTime)
import System.IO (hFlush, stdout)
import System.Mem (performGC)

import Imp.DSL (Imp)
import Imp.Inference (intervalProbability, intervalProbabilityApprox)
import Imp.Inference.Optimize (optimizeProbability)
import Imp.Inference.Symbolic (intervalProbabilitySymbolic)

-- | A single benchmark case.
data BenchCase where
  BenchCase :: Ord a => Int -> Imp g a -> (a -> Bool) -> BenchCase

-- | Median timing across the four inference methods, in seconds.
data BenchResult = BenchResult
  { brN        :: Int
  , brExact    :: Maybe Double
  , brSymbolic :: Maybe Double
  , brApprox   :: Double
  , brGrad     :: Double
  }

data BenchConfig = BenchConfig
  { bcReps         :: Int        -- ^ median over this many reps
  , bcGradSteps    :: Int        -- ^ gradient-descent step count
  , bcGradLR       :: Double     -- ^ gradient-descent learning rate
  , bcMaxExactN    :: Maybe Int  -- ^ skip exact enumeration when @N > this@
  , bcMaxSymbolicN :: Maybe Int  -- ^ skip symbolic when @N > this@
  , bcGradBounds   :: Bool       -- ^ run gradient in both directions
  , bcGcBetween    :: Bool       -- ^ 'performGC' before each timed rep
  }

-- | Defaults tuned for the Ellsberg-style benches (fast individual reps,
--   high-dimensional gradient landscapes).  Slower benches (e.g. Robot)
--   override 'bcGradSteps', 'bcGradLR', and 'bcGradBounds' to suit.
defaultConfig :: BenchConfig
defaultConfig = BenchConfig
  { bcReps         = 10
  , bcGradSteps    = 30
  , bcGradLR       = 0.3
  , bcMaxExactN    = Just 20
  , bcMaxSymbolicN = Just 20
  , bcGradBounds   = True
  , bcGcBetween    = True
  }

-- ---------------------------------------------------------------------------
-- Pipeline
-- ---------------------------------------------------------------------------

-- | Run every case, emitting CSV (header + per-case rows) to stdout in real-time.
runAll :: BenchConfig -> [BenchCase] -> IO ()
runAll cfg cs = do
  putStrLn csvHeader
  hFlush stdout
  mapM_ (\c -> runBench cfg c >>= emitCsvRow) cs

-- | Run a single benchmark case, timing each of the four inference
--   methods and returning their medians.
runBench :: BenchConfig -> BenchCase -> IO BenchResult
runBench cfg (BenchCase n prog predicate) = do
  let skipExact    = maybe False (n >) (bcMaxExactN cfg)
      skipSymbolic = maybe False (n >) (bcMaxSymbolicN cfg)
      steps        = bcGradSteps cfg
      lr           = bcGradLR cfg
  tExact    <- if skipExact
                 then return Nothing
                 else Just <$> timeFresh cfg (\() -> intervalProbability prog predicate)
  tSymbolic <- if skipSymbolic
                 then return Nothing
                 else Just <$> timeFresh cfg (\() -> intervalProbabilitySymbolic prog predicate)
  tApprox   <- timeFresh cfg (\() -> intervalProbabilityApprox prog predicate)
  tGrad     <- timeFresh cfg $ \() ->
    if bcGradBounds cfg
      then gradientBounds steps lr prog predicate
      else let (_, v) = optimizeProbability prog predicate steps lr
           in (v, v)
  return (BenchResult n tExact tSymbolic tApprox tGrad)

-- | Gradient-descent bounds.
gradientBounds :: Ord a => Int -> Double -> Imp g a -> (a -> Bool) -> (Double, Double)
gradientBounds nSteps lr prog predicate =
  let (_, pMax) = optimizeProbability prog predicate nSteps   lr
      (_, pMin) = optimizeProbability prog predicate nSteps (-lr)
  in (pMin, pMax)

-- ---------------------------------------------------------------------------
-- Timing
-- ---------------------------------------------------------------------------

-- | Median wall-clock time (in seconds) of @bcReps@ fresh evaluations of
--   the supplied nullary function.
timeFresh :: BenchConfig -> (() -> (Double, Double)) -> IO Double
timeFresh cfg compute = do
  ts <- replicateM (bcReps cfg) $ do
    when (bcGcBetween cfg) performGC
    (_, t) <- timeIt $
      let (a, b) = compute () in a `seq` b `seq` return (a, b)
    return t
  return (median ts)

timeIt :: IO a -> IO (a, Double)
timeIt act = do
  t0 <- getCurrentTime
  x  <- act
  t1 <- getCurrentTime
  return (x, realToFrac (diffUTCTime t1 t0))

median :: [Double] -> Double
median xs =
  let s = sort xs
      m = length s
  in if odd m
       then s !! (m `div` 2)
       else 0.5 * (s !! (m `div` 2 - 1) + s !! (m `div` 2))

-- ---------------------------------------------------------------------------
-- CSV output
-- ---------------------------------------------------------------------------

-- | CSV column header: @N,method,time_s@.
csvHeader :: String
csvHeader = "N,method,time_s"

-- | Four CSV rows (one per method) for a single result.
emitCsvRow :: BenchResult -> IO ()
emitCsvRow r = do
  let n = brN r
  putStrLn (show n ++ ",exact,"    ++ maybe "NaN" show (brExact r))
  putStrLn (show n ++ ",symbolic," ++ maybe "NaN" show (brSymbolic r))
  putStrLn (show n ++ ",interval," ++ show (brApprox r))
  putStrLn (show n ++ ",gradient," ++ show (brGrad r))
  hFlush stdout