eproc

Anytime-valid sequential testing and confidence sequences.
git clone git://git.ppad.tech/eproc.git
Log | Files | Refs | README | LICENSE

Bounded.hs (13755B)


      1 {-# OPTIONS_HADDOCK prune #-}
      2 {-# LANGUAGE BangPatterns #-}
      3 {-# LANGUAGE RecordWildCards #-}
      4 
      5 -- |
      6 -- Module: Numeric.Eproc.Bounded
      7 -- Copyright: (c) 2026 Jared Tobin
      8 -- License: MIT
      9 -- Maintainer: Jared Tobin <jared@ppad.tech>
     10 --
     11 -- Two-sided bounded-mean anytime-valid test.
     12 --
     13 -- For samples @x_t@ in @[lo, hi]@, tests
     14 --
     15 --     @H_0: E[x_t | F_{t-1}] = m   for all t@
     16 --
     17 -- against the negation. Here @F_{t-1}@ is the filtration generated
     18 -- by everything observed strictly before time @t@; the conditional
     19 -- form is what anytime validity actually requires. For i.i.d.
     20 -- samples this reduces to the usual marginal statement
     21 -- @E[x] = m@; for adaptively-collected or otherwise non-i.i.d.
     22 -- streams the conditional statement is the right thing to think
     23 -- about.
     24 --
     25 -- Internally two one-sided e-processes are run in parallel: a
     26 -- /positive-direction/ process @K^+_t@ betting against the
     27 -- alternative @E[x_t | F_{t-1}] > m@ (using centred observations
     28 -- @z = x - m@), and a /negative-direction/ process @K^-_t@ betting
     29 -- against @E[x_t | F_{t-1}] < m@ (using @-z@). Each maintains its
     30 -- own log-wealth and bettor state.
     31 --
     32 -- The two sides are combined via the /hedged capital process/ of
     33 -- Waudby-Smith & Ramdas (2024) §4: their average
     34 -- @K_t = (K^+_t + K^-_t) \/ 2@ is itself an e-process (convex
     35 -- combinations preserve the supermartingale property), with
     36 -- @E[K_0] = 1@. By Ville's inequality
     37 -- @P(sup_t K_t >= 1 \/ alpha) <= alpha@, so the test rejects when
     38 -- the supremum of @K^+_t + K^-_t@ has ever crossed @2 \/ alpha@.
     39 --
     40 -- This is strictly more powerful than the naive Bonferroni union
     41 -- (reject when @max(K^+_t, K^-_t) >= 2 \/ alpha@): the convex-hedge
     42 -- rejection region contains Bonferroni's (since
     43 -- @K^+ + K^- >= max(K^+, K^-)@), with the same alpha guarantee.
     44 -- For one-sided alternatives the gap is small (the losing-direction
     45 -- bettor stays near @1@); for genuinely two-sided alternatives it
     46 -- can be substantial.
     47 --
     48 -- The test is /anytime-valid/: under @H_0@ the wealth process is a
     49 -- nonnegative supermartingale, so by Ville's inequality the
     50 -- probability of /ever/ crossing the threshold is at most @alpha@,
     51 -- regardless of when the user decides to stop streaming samples.
     52 -- Rejection is /latched/ in the running state -- once a side has
     53 -- crossed threshold, 'decide' continues to return 'Reject' even if
     54 -- the current log-wealth has since dropped back below threshold.
     55 --
     56 -- == Example
     57 --
     58 -- Test @H_0: E[x] = 0.5@ for @x@ in @[0, 1]@ at level @alpha = 1e-3@
     59 -- against a stream with empirical mean @0.8@:
     60 --
     61 -- >>> let Right cfg = config 0.5 0.0 1.0 1.0e-3 Newton
     62 -- >>> let xs  = concat (replicate 30 [1, 1, 0, 1, 1, 0, 1, 1, 1, 1])
     63 -- >>> decide cfg (foldl' (update cfg) (initial cfg) xs)
     64 -- Reject
     65 
     66 module Numeric.Eproc.Bounded (
     67   -- * Test configuration and state
     68     Config
     69   , State
     70   , Verdict(..)
     71   , ConfigError(..)
     72 
     73   -- * Bettor strategies
     74   , Bettor(..)
     75 
     76   -- * Construction
     77   , config
     78   , initial
     79 
     80   -- * Streaming
     81   , update
     82   , decide
     83 
     84   -- * Inspection
     85   , log_wealth
     86   , log_wealth_sup
     87   , log_evalue
     88   , log_evalue_sup
     89   , p_value
     90   , samples
     91   ) where
     92 
     93 import GHC.Float (log1p)
     94 import Numeric.Eproc.Common (
     95     Bettor(..), Verdict(..), ConfigError(..)
     96   , BetState, init_bet, bet_lambda, step_bet
     97   , finite, log_sum_exp, log2_dbl
     98   )
     99 
    100 -- types ----------------------------------------------------------------------
    101 
    102 -- here, the centred observation @z_t@ referenced in
    103 -- "Numeric.Eproc.Common" is @x_t - m@; the per-direction safe-bet
    104 -- ceilings @lambda_max@ are derived from the sample bounds (see
    105 -- 'config').
    106 
    107 -- | Bounded-mean test configuration. Build with 'config'.
    108 --
    109 --   Carries the bettor strategy, the null mean, the significance
    110 --   level, the precomputed convex-hedge log-wealth threshold, and
    111 --   the per-direction safe-bet ceilings (see 'config' for how the
    112 --   latter are derived from the sample bounds).
    113 data Config = Config {
    114     -- ^ bettor strategy
    115     cfg_bettor      :: !Bettor
    116     -- ^ positive-direction safe-bet ceiling
    117   , cfg_lam_max_pos :: {-# UNPACK #-} !Double
    118     -- ^ negative-direction safe-bet ceiling
    119   , cfg_lam_max_neg :: {-# UNPACK #-} !Double
    120     -- ^ null mean @m@
    121   , cfg_null_mean   :: {-# UNPACK #-} !Double
    122     -- ^ significance level @alpha@
    123   , cfg_alpha       :: {-# UNPACK #-} !Double
    124     -- ^ rejection threshold @log(2 \/ alpha)@
    125   , cfg_log_thresh  :: {-# UNPACK #-} !Double
    126   }
    127 
    128 -- | Streaming test state. Construct with 'initial' and fold
    129 --   observations through 'update'.
    130 --
    131 --   The two log-wealth fields track the running log-wealth of the
    132 --   positive- and negative-direction e-processes separately; the
    133 --   /sup log-sum/ field latches the supremum so far of
    134 --   @log(K^+_t + K^-_t)@, which is the test statistic the
    135 --   convex-hedge construction actually monitors. The per-direction
    136 --   bettor states carry whatever the chosen 'Bettor' needs (running
    137 --   sums, current bet, etc.).
    138 data State = State {
    139     st_n           :: {-# UNPACK #-} !Int     -- ^ sample count
    140   , st_log_w_pos   :: {-# UNPACK #-} !Double  -- ^ log-wealth, pos
    141   , st_log_w_neg   :: {-# UNPACK #-} !Double  -- ^ log-wealth, neg
    142   , st_sup_log_sum :: {-# UNPACK #-} !Double  -- ^ sup log(K^+ + K^-)
    143   , st_bet_pos     :: !BetState               -- ^ bettor state, pos
    144   , st_bet_neg     :: !BetState               -- ^ bettor state, neg
    145   }
    146 
    147 -- construction ---------------------------------------------------------------
    148 
    149 -- | Build a 'Config' for the bounded-mean test.
    150 --
    151 --   Each per-direction safe-bet ceiling @lambda_max@ is set so that
    152 --   the wealth factor stays nonnegative for every admissible
    153 --   observation:
    154 --
    155 --   * The positive-direction factor is @1 + lambda_p * (x - m)@.
    156 --     Since @x@ can dip to @lo@, @x - m@ can reach @lo - m@ (the
    157 --     most negative value), so we need
    158 --     @lambda_p <= 1 \/ (m - lo)@. The ceiling stored is half this
    159 --     to leave numerical margin -- the WSR safety recommendation.
    160 --
    161 --   * The negative-direction factor is @1 - lambda_n * (x - m)@.
    162 --     Since @x@ can rise to @hi@, @x - m@ can reach @hi - m@, so we
    163 --     need @lambda_n <= 1 \/ (hi - m)@; again the ceiling is set to
    164 --     half this.
    165 --
    166 --   The log-wealth rejection threshold is precomputed as
    167 --   @log(2 \/ alpha)@; the 2 reflects that the convex-hedge test
    168 --   monitors the sum @K^+_t + K^-_t@, whose initial value is @2@
    169 --   (each side starts at @K = 1@).
    170 --
    171 --   Returns 'Left' with a 'ConfigError' on inputs that would leave
    172 --   the mathematical regime: any of @m@, @lo@, @hi@, @alpha@
    173 --   non-finite (NaN or infinite); @alpha@ outside @(0, 1)@;
    174 --   @lo >= hi@; or @m@ outside the open interval @(lo, hi)@
    175 --   (strict, to avoid the safe-bet ceilings dividing by zero).
    176 --
    177 --   >>> let Right cfg = config 0.5 0.0 1.0 1.0e-3 Newton
    178 config
    179   :: Double  -- ^ null mean @m@
    180   -> Double  -- ^ sample lower bound @lo@
    181   -> Double  -- ^ sample upper bound @hi@
    182   -> Double  -- ^ significance level @alpha@
    183   -> Bettor  -- ^ bettor strategy
    184   -> Either ConfigError Config
    185 config !m !lo !hi !alpha !b
    186   | not (finite alpha && alpha > 0 && alpha < 1) =
    187       Left (InvalidAlpha alpha)
    188   | not (finite lo && finite hi && lo < hi) =
    189       Left (InvalidBounds lo hi)
    190   | not (finite m && lo < m && m < hi) =
    191       Left (InvalidNullMean m lo hi)
    192   | otherwise = Right Config {
    193         cfg_bettor      = b
    194       , cfg_lam_max_pos = 0.5 / (m - lo)
    195       , cfg_lam_max_neg = 0.5 / (hi - m)
    196       , cfg_null_mean   = m
    197       , cfg_alpha       = alpha
    198       , cfg_log_thresh  = log (2 / alpha)
    199       }
    200 {-# INLINE config #-}
    201 
    202 -- | The initial 'State' for a fresh streaming test.
    203 --
    204 --   Both per-direction log-wealths start at @0@ (i.e., @K = 1@);
    205 --   the sup-log-sum starts at @log 2@ (since @K^+_0 + K^-_0 = 2@);
    206 --   both bettors start in the per-strategy initial state
    207 --   appropriate for the 'Bettor' chosen in the 'Config'.
    208 --
    209 --   >>> let s0 = initial cfg
    210 initial :: Config -> State
    211 initial Config{..} =
    212   let !s0 = init_bet cfg_bettor
    213   in  State {
    214         st_n           = 0
    215       , st_log_w_pos   = 0
    216       , st_log_w_neg   = 0
    217       , st_sup_log_sum = log2_dbl
    218       , st_bet_pos     = s0
    219       , st_bet_neg     = s0
    220       }
    221 {-# INLINE initial #-}
    222 
    223 -- streaming ------------------------------------------------------------------
    224 
    225 -- | Fold one observation into the running 'State'.
    226 --
    227 --   Computes the centred observation @z = x - m@, queries the two
    228 --   directional bettors for their predictable bets, accumulates
    229 --   per-direction log-wealth via
    230 --
    231 --       @log_w' = log_w + log (1 + lambda * z)@
    232 --
    233 --   (with the symmetric @-lambda@ for the negative direction), then
    234 --   updates the running supremum of @log(K^+ + K^-)@ via
    235 --   log-sum-exp and steps the bettor states given the newly
    236 --   observed @z@.
    237 --
    238 --   /Precondition/: @x@ must lie in the @[lo, hi]@ interval given
    239 --   to 'config'. The type-I error guarantee of the test depends on
    240 --   this. Out-of-range observations can drive the wealth factor
    241 --   negative, taking the construction out of the supermartingale
    242 --   regime entirely; the function does not check for this.
    243 --
    244 --   >>> let s1 = update cfg s0 0.7
    245 update :: Config -> State -> Double -> State
    246 update Config{..} State{..} !x =
    247   let !z       = x - cfg_null_mean
    248       !lam_p   = bet_lambda cfg_bettor cfg_lam_max_pos st_bet_pos
    249       !lam_n   = bet_lambda cfg_bettor cfg_lam_max_neg st_bet_neg
    250       !logw_p  = st_log_w_pos + log1p (lam_p * z)
    251       !logw_n  = st_log_w_neg + log1p (negate lam_n * z)
    252       -- Skip 'log_sum_exp' when the cheap upper bound
    253       --   log_sum_exp a b <= max a b + log 2
    254       -- already sits at or below the running max: no update can
    255       -- move it. Under H_0 (calibration) this is the common case.
    256       !cheap_ub = max logw_p logw_n + log2_dbl
    257       !sup_sum
    258         | cheap_ub <= st_sup_log_sum = st_sup_log_sum
    259         | otherwise                  =
    260             max st_sup_log_sum (log_sum_exp logw_p logw_n)
    261       !sp      = step_bet cfg_bettor cfg_lam_max_pos st_bet_pos z
    262       !sn      = step_bet cfg_bettor cfg_lam_max_neg st_bet_neg (negate z)
    263   in  State (st_n + 1) logw_p logw_n sup_sum sp sn
    264 {-# INLINE update #-}
    265 
    266 -- | Compute the current 'Verdict' from the running 'State'.
    267 --
    268 --   'Reject' iff the supremum-so-far of @log(K^+_t + K^-_t)@ has
    269 --   ever crossed the threshold @log(2 \/ alpha)@ — equivalently,
    270 --   the convex-hedge e-process @(K^+ + K^-) \/ 2@ has exceeded
    271 --   @1 \/ alpha@ at some point in the stream so far. Under @H_0@,
    272 --   by Ville's inequality, the probability of this ever happening
    273 --   is at most @alpha@ -- and crucially this bound holds at /every/
    274 --   sample size simultaneously, so the user is free to peek at the
    275 --   verdict as often as they like and stop on the first 'Reject'.
    276 --
    277 --   >>> decide cfg s0
    278 --   Continue
    279 decide :: Config -> State -> Verdict
    280 decide Config{..} State{..}
    281   | st_sup_log_sum >= cfg_log_thresh = Reject
    282   | otherwise                        = Continue
    283 {-# INLINE decide #-}
    284 
    285 -- inspection -----------------------------------------------------------------
    286 
    287 -- | The current @log(K^+_t + K^-_t)@ -- the running log-wealth of
    288 --   the convex-hedge combination at the present sample count.
    289 --
    290 --   Unlike 'log_wealth_sup' this is not monotone: adverse
    291 --   observations decrease it. It is bounded above by
    292 --   'log_wealth_sup', which is what 'decide' tests against the
    293 --   rejection threshold.
    294 --
    295 --   Starts at @log 2@ (since @K^+_0 + K^-_0 = 2@).
    296 --
    297 --   >>> log_wealth s0
    298 --   0.6931471805599453
    299 log_wealth :: State -> Double
    300 log_wealth State{..} = log_sum_exp st_log_w_pos st_log_w_neg
    301 {-# INLINE log_wealth #-}
    302 
    303 -- | The supremum-so-far of @log(K^+_t + K^-_t)@, taken across all
    304 --   sample counts up to the current one. This is the test statistic
    305 --   the convex-hedge construction actually monitors: it is monotone
    306 --   nondecreasing in the sample count, and 'decide' rejects exactly
    307 --   when it crosses @log(2 \/ alpha)@.
    308 --
    309 --   Starts at @log 2@ (since @K^+_0 + K^-_0 = 2@).
    310 --
    311 --   >>> log_wealth_sup s0
    312 --   0.6931471805599453
    313 log_wealth_sup :: State -> Double
    314 log_wealth_sup State{..} = st_sup_log_sum
    315 {-# INLINE log_wealth_sup #-}
    316 
    317 -- | The current log e-value of the convex-hedge e-process: the log
    318 --   of @(K^+_t + K^-_t) \/ 2@, i.e. 'log_wealth' minus @log 2@.
    319 --
    320 --   Unlike 'log_wealth', this is normalized so that a fresh state
    321 --   sits at @0@ (e-value @1@): it is directly comparable across
    322 --   test modules regardless of their internal hedging, and is the
    323 --   form to use when convex-combining several e-processes. Not
    324 --   monotone; bounded above by 'log_evalue_sup'.
    325 --
    326 --   >>> log_evalue s0
    327 --   0.0
    328 log_evalue :: State -> Double
    329 log_evalue s = log_wealth s - log2_dbl
    330 {-# INLINE log_evalue #-}
    331 
    332 -- | The supremum-so-far of the log e-value: 'log_wealth_sup' minus
    333 --   @log 2@. Monotone nondecreasing, starting at @0@; 'decide'
    334 --   rejects exactly when it crosses @log(1 \/ alpha)@.
    335 --
    336 --   >>> log_evalue_sup s0
    337 --   0.0
    338 log_evalue_sup :: State -> Double
    339 log_evalue_sup s = log_wealth_sup s - log2_dbl
    340 {-# INLINE log_evalue_sup #-}
    341 
    342 -- | The anytime-valid p-value: the reciprocal of the largest
    343 --   e-value attained so far, @min 1 (exp (negate (log_evalue_sup
    344 --   s)))@.
    345 --
    346 --   Monotone nonincreasing in the sample count, and valid under
    347 --   optional stopping: under @H_0@,
    348 --   @P(exists t: p_t <= alpha) <= alpha@ for every @alpha@
    349 --   simultaneously. 'decide' returns 'Reject' exactly when this
    350 --   value has reached the configured @alpha@ or below.
    351 --
    352 --   >>> p_value s0
    353 --   1.0
    354 p_value :: State -> Double
    355 p_value s = min 1 (exp (negate (log_evalue_sup s)))
    356 {-# INLINE p_value #-}
    357 
    358 -- | The number of samples consumed so far.
    359 --
    360 --   >>> samples s0
    361 --   0
    362 samples :: State -> Int
    363 samples = st_n
    364 {-# INLINE samples #-}