Bernoulli.hs

{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}

-- |
-- Module: Numeric.Eproc.Bernoulli
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- One-sided Bernoulli rate anytime-valid test. See
-- "Numeric.Eproc.Bernoulli.TwoSided" for the two-sided companion
-- (used for the sign test at @p_0 = 1\/2@, among other things).
--
-- For samples @x_t@ in @{0, 1}@, tests
--
--     @H_0: E[x_t | F_{t-1}] <= p_0   for all t@
--
-- against @H_1: E[x_t | F_{t-1}] > p_0@ (at some @t@). Here
-- @F_{t-1}@ is the filtration generated by everything observed
-- strictly before time @t@; the conditional form is what anytime
-- validity actually requires. For i.i.d. samples this reduces to
-- the usual marginal statement @E[x] <= p_0@.
--
-- A single wealth process is run:
--
--     @W_n = prod_{i=1..n} (1 + lambda_i * (x_i - p_0))@
--
-- where each per-step bet @lambda_i@ is chosen predictably (from
-- data observed strictly before step @i@) and clipped to
-- @[0, lambda_max]@ so that the wealth factor stays nonnegative for
-- every admissible observation. Under @H_0@ the wealth process is
-- a nonnegative supermartingale, so by Ville's inequality the
-- probability of @W_n@ ever crossing @1 \/ alpha@ is at most
-- @alpha@, regardless of when the user decides to stop streaming
-- samples. Rejection is /latched/ in the running state: once the
-- wealth has crossed threshold, 'decide' continues to return
-- 'Reject' even if subsequent observations drive the current
-- wealth back below threshold.
--
-- The alternative here is one-sided, so a single wealth process
-- suffices and no Bonferroni or hedge adjustment is needed -- the
-- rejection threshold is @log(1 \/ alpha)@.
--
-- == Example
--
-- Test @H_0: E[x] <= 0.05@ at level @alpha = 1e-3@ against a stream
-- with empirical rate @~0.5@:
--
-- >>> let Right cfg = config 0.05 1.0e-3 Newton
-- >>> let xs  = take 200 (cycle [True, False])
-- >>> decide cfg (foldl' (update cfg) (initial cfg) xs)
-- Reject

module Numeric.Eproc.Bernoulli (
  -- * Test configuration and state
    Config
  , State
  , Verdict(..)
  , ConfigError(..)

  -- * Bettor strategies
  , Bettor(..)

  -- * Construction
  , config
  , initial

  -- * Streaming
  , update
  , decide

  -- * Inspection
  , log_wealth
  , samples
  ) where

import GHC.Float (log1p)
import Numeric.Eproc.Common (
    Bettor(..), Verdict(..), ConfigError(..)
  , BetState, init_bet, bet_lambda, step_bet
  , finite
  )

-- types ----------------------------------------------------------------------

-- here, the centred observation @z_t@ referenced in
-- "Numeric.Eproc.Common" is @x_t - p_0@; the safe-bet ceiling
-- @lambda_max@ is derived from @p_0@ (see 'config').

-- | Bernoulli rate test configuration. Build with 'config'.
--
--   Carries the bettor strategy, the baseline rate, the significance
--   level, the precomputed log-wealth rejection threshold, and the
--   safe-bet ceiling derived from @p_0@.
data Config = Config {
    -- ^ bettor strategy
    cfg_bettor     :: !Bettor
    -- ^ safe-bet ceiling
  , cfg_lam_max    :: {-# UNPACK #-} !Double
    -- ^ baseline rate @p_0@
  , cfg_p0         :: {-# UNPACK #-} !Double
    -- ^ significance level @alpha@
  , cfg_alpha      :: {-# UNPACK #-} !Double
    -- ^ rejection threshold @log(1 \/ alpha)@
  , cfg_log_thresh :: {-# UNPACK #-} !Double
  }

-- | Streaming test state. Construct with 'initial' and fold
--   observations through 'update'.
--
--   Carries the sample count, current and supremum-so-far running
--   log-wealth, and whatever per-step state the chosen 'Bettor'
--   needs. The supremum field is what 'decide' tests against the
--   rejection threshold; this is the supremum-style event Ville's
--   inequality actually bounds.
data State = State {
    st_n         :: {-# UNPACK #-} !Int     -- ^ sample count
  , st_log_w     :: {-# UNPACK #-} !Double  -- ^ running log-wealth
  , st_max_log_w :: {-# UNPACK #-} !Double  -- ^ sup log-wealth so far
  , st_bet       :: !BetState               -- ^ bettor state
  }

-- construction ---------------------------------------------------------------

-- | Build a 'Config' for the Bernoulli rate test.
--
--   The safe-bet ceiling @lambda_max@ is set so that the wealth
--   factor @1 + lambda * (x - p_0)@ stays nonnegative for both
--   @x = 0@ and @x = 1@. The binding constraint is @x = 0@, which
--   requires @lambda <= 1 \/ p_0@; the ceiling stored is half this
--   to leave numerical margin -- the WSR safety recommendation.
--
--   Returns 'Left' with a 'ConfigError' on inputs that would leave
--   the mathematical regime: either of @p_0@ or @alpha@ non-finite
--   (NaN or infinite); @p_0@ outside @(0, 1)@ (the degenerate case
--   @p_0 = 0@ would make @lambda_max@ infinite, and @p_0 = 1@
--   leaves no room for an alternative); or @alpha@ outside
--   @(0, 1)@.
--
--   >>> let Right cfg = config 0.05 1.0e-3 Newton
config
  :: Double  -- ^ baseline rate @p_0@, in @(0, 1)@
  -> Double  -- ^ significance level @alpha@, in @(0, 1)@
  -> Bettor  -- ^ bettor strategy
  -> Either ConfigError Config
config !p0 !alpha !b
  | not (finite p0 && p0 > 0 && p0 < 1) =
      Left (InvalidBaselineRate p0)
  | not (finite alpha && alpha > 0 && alpha < 1) =
      Left (InvalidAlpha alpha)
  | otherwise = Right Config {
        cfg_bettor     = b
      , cfg_lam_max    = 0.5 / p0
      , cfg_p0         = p0
      , cfg_alpha      = alpha
      , cfg_log_thresh = log (1 / alpha)
      }
{-# INLINE config #-}

-- | The initial 'State' for a fresh streaming test.
--
--   Both log-wealth fields start at @0@ (i.e., wealth @1@) and the
--   bettor starts in the per-strategy initial state appropriate
--   for the 'Bettor' chosen in the 'Config'.
--
--   >>> let s0 = initial cfg
initial :: Config -> State
initial Config{..} = State {
    st_n         = 0
  , st_log_w     = 0
  , st_max_log_w = 0
  , st_bet       = init_bet cfg_bettor
  }
{-# INLINE initial #-}

-- streaming ------------------------------------------------------------------

-- | Fold one observation into the running 'State'.
--
--   @True@ means @x_t = 1@ (the event of interest occurred -- e.g.,
--   two readings diverged); @False@ means @x_t = 0@ (they matched).
--   The caller decides what \"matched\" means at the application
--   level.
--
--   Computes the centred observation @z = x - p_0@, queries the
--   bettor for its predictable bet, accumulates log-wealth via
--
--       @log_w' = log_w + log (1 + lambda * z)@
--
--   updates the running supremum log-wealth, then steps the bettor
--   state given the newly observed @z@.
--
--   /Precondition/: @True@ and @False@ both /must/ be admissible
--   under the test (this holds vacuously for the @{0, 1}@ support).
--   The function is total.
--
--   >>> let s1 = update cfg s0 True
update :: Config -> State -> Bool -> State
update Config{..} State{..} !x =
  let !xd     = if x then 1 else 0
      !z      = xd - cfg_p0
      !lam    = bet_lambda cfg_bettor cfg_lam_max st_bet
      !logw'  = st_log_w + log1p (lam * z)
      !maxw'  = max st_max_log_w logw'
      !s'     = step_bet cfg_bettor cfg_lam_max st_bet z
  in  State (st_n + 1) logw' maxw' s'
{-# INLINE update #-}

-- | Compute the current 'Verdict' from the running 'State'.
--
--   'Reject' iff log-wealth has /ever/ crossed the threshold
--   @log(1 \/ alpha)@; equivalently, wealth has exceeded
--   @1 \/ alpha@ at some point in the stream so far. Under @H_0@,
--   by Ville's inequality, the probability of this ever happening
--   is at most @alpha@ -- and crucially this bound holds at /every/
--   sample size simultaneously, so the user is free to peek at the
--   verdict as often as they like and stop on the first 'Reject'.
--
--   >>> decide cfg s0
--   Continue
decide :: Config -> State -> Verdict
decide Config{..} State{..}
  | st_max_log_w >= cfg_log_thresh = Reject
  | otherwise                      = Continue
{-# INLINE decide #-}

-- inspection -----------------------------------------------------------------

-- | The supremum-so-far log-wealth, across all sample counts up to
--   the current one.
--
--   This is the natural \"test statistic\": it is monotone
--   nondecreasing in the sample count, and 'decide' rejects exactly
--   when it crosses @log(1 \/ alpha)@.
--
--   >>> log_wealth s0
--   0.0
log_wealth :: State -> Double
log_wealth = st_max_log_w
{-# INLINE log_wealth #-}

-- | The number of samples consumed so far.
--
--   >>> samples s0
--   0
samples :: State -> Int
samples = st_n
{-# INLINE samples #-}