eproc

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

Common.hs (9239B)


      1 {-# OPTIONS_HADDOCK prune #-}
      2 {-# LANGUAGE BangPatterns #-}
      3 
      4 -- |
      5 -- Module: Numeric.Eproc.Common
      6 -- Copyright: (c) 2026 Jared Tobin
      7 -- License: MIT
      8 -- Maintainer: Jared Tobin <jared@ppad.tech>
      9 --
     10 -- Shared vocabulary for the eproc tests: the predictable bettor
     11 -- strategies, the test verdict type, and the configuration-error
     12 -- type. Re-exported from each test module
     13 -- ("Numeric.Eproc.Bounded", "Numeric.Eproc.Paired",
     14 -- "Numeric.Eproc.Bernoulli"); import this module directly only if
     15 -- you need the types without picking a particular test.
     16 --
     17 -- The 'BetState' type and its helpers are internal to the library:
     18 -- they are exposed here so that 'Numeric.Eproc.Bounded' and
     19 -- 'Numeric.Eproc.Bernoulli' can share one implementation, not for
     20 -- direct use.
     21 
     22 module Numeric.Eproc.Common (
     23     Bettor(..)
     24   , Verdict(..)
     25   , ConfigError(..)
     26 
     27   -- * Internal: shared bettor state
     28   , BetState(..)
     29   , init_bet
     30   , bet_lambda
     31   , step_bet
     32 
     33   -- * Internal: helpers
     34   , finite
     35   , log_sum_exp
     36   , log2_dbl
     37   ) where
     38 
     39 import GHC.Float (log1p)
     40 
     41 -- | A predictable bettor.
     42 --
     43 --   A bettor describes how, given the history of centred
     44 --   observations @z_t@ (each test module specifies its own centring;
     45 --   see the per-module documentation), the next predictable bet
     46 --   @lambda_t@ is chosen. Predictability -- that is, @lambda_t@
     47 --   depends only on data observed strictly before step @t@ -- is
     48 --   what makes the resulting wealth process a nonnegative
     49 --   supermartingale under @H_0@.
     50 --
     51 --   All three bettors enforce a safe-bet ceiling @lambda_max@
     52 --   derived from the test's admissible-observation range by clipping
     53 --   @lambda@ to @[0, lambda_max]@; this keeps the per-step wealth
     54 --   factor nonnegative.
     55 --
     56 --   * 'Fixed' bets the supplied constant @lambda@ (clipped to
     57 --     @[0, lambda_max]@). The wager does not respond to observed
     58 --     data; this strategy is useful only as a baseline.
     59 --
     60 --   * 'Adaptive' is the aGRAPA (approximate growth-rate adaptive
     61 --     predictable plug-in) bettor of Waudby-Smith & Ramdas (2024).
     62 --     It tracks the empirical mean @mu@ and variance @sigma^2@ of
     63 --     centred observations and bets the Kelly-optimal plug-in
     64 --     @lambda* = mu \/ (sigma^2 + mu^2)@ clipped to
     65 --     @[0, lambda_max]@. Fast to compute and competitive in
     66 --     practice.
     67 --
     68 --   * 'Newton' is the online Newton step (ONS) bettor of
     69 --     Waudby-Smith & Ramdas (2024, Algorithm 2). The per-step
     70 --     log-wealth loss @-log(1 + lambda * z)@ is convex in @lambda@;
     71 --     ONS performs one Newton step per observation, accumulating
     72 --     squared gradients to scale the update by a fixed learning
     73 --     rate @2 \/ (2 - log 3)@. Achieves logarithmic regret against
     74 --     the best constant bet in hindsight and is in practice the
     75 --     strongest of the three bettors under most signal regimes.
     76 --
     77 --     One deliberate deviation from WSR: Algorithm 2 seeds the
     78 --     squared-gradient accumulator at @1@, which presumes
     79 --     observations scaled to @[0, 1]@. On raw-scale data that
     80 --     constant is dimensionally wrong -- negligible when
     81 --     @z^2 >> 1@, paralysing when @z^2 << 1@ -- so the accumulator
     82 --     here is instead seeded near zero, making the update
     83 --     scale-adaptive. The trade is bold early play: the first
     84 --     nonzero observation typically drives the bet straight to
     85 --     the @lambda_max@ ceiling, annealing back toward the Kelly
     86 --     point as gradients accumulate. Validity is unaffected --
     87 --     predictability and clipping are all it needs -- and regret
     88 --     stays logarithmic with a somewhat larger constant. The
     89 --     visible effect is higher-variance early wealth: a supremum
     90 --     modestly above its floor is expected even under @H_0@.
     91 data Bettor =
     92     Fixed {-# UNPACK #-} !Double
     93   | Adaptive
     94   | Newton
     95   deriving (Eq, Show)
     96 
     97 -- | Test outcome at the current sample count.
     98 --
     99 --   'Reject' means the wealth process has /ever/ crossed the
    100 --   rejection threshold, so @H_0@ is rejected at level @alpha@.
    101 --   Once a state has rejected it stays rejected, even if subsequent
    102 --   observations drive the current wealth back below threshold;
    103 --   this is the supremum-style guarantee that Ville's inequality
    104 --   actually delivers. 'Continue' means there is not yet enough
    105 --   evidence; collect more samples (or stop and report no
    106 --   rejection -- the type-I error guarantee holds for /any/
    107 --   stopping rule).
    108 data Verdict =
    109     Reject
    110   | Continue
    111   deriving (Eq, Show)
    112 
    113 -- | Reasons that a test-configuration smart constructor can reject
    114 --   its inputs. Returned by 'Numeric.Eproc.Bounded.config',
    115 --   'Numeric.Eproc.Bernoulli.config',
    116 --   'Numeric.Eproc.Paired.config',
    117 --   'Numeric.Eproc.Mixture.config', and
    118 --   'Numeric.Eproc.ConfSeq.config'.
    119 data ConfigError =
    120     -- | significance level outside @(0, 1)@
    121     InvalidAlpha {-# UNPACK #-} !Double
    122     -- | sample bounds violate @lo < hi@
    123   | InvalidBounds {-# UNPACK #-} !Double {-# UNPACK #-} !Double
    124     -- | null mean outside @(lo, hi)@ (strict, to avoid div-by-zero
    125     --   in the safe-bet ceilings)
    126   | InvalidNullMean
    127       {-# UNPACK #-} !Double  -- m
    128       {-# UNPACK #-} !Double  -- lo
    129       {-# UNPACK #-} !Double  -- hi
    130     -- | baseline rate outside @(0, 1)@
    131   | InvalidBaselineRate {-# UNPACK #-} !Double
    132     -- | component count not positive
    133   | InvalidComponentCount {-# UNPACK #-} !Int
    134     -- | grid size below @1@
    135   | InvalidGridSize {-# UNPACK #-} !Int
    136   deriving (Eq, Show)
    137 
    138 -- | True iff the argument is a finite IEEE-754 double (not NaN, not
    139 --   @+\/-Infinity@). Used by the @config@ smart constructors to keep
    140 --   the bounded-random-variable promise honest.
    141 finite :: Double -> Bool
    142 finite x = not (isNaN x) && not (isInfinite x)
    143 {-# INLINE finite #-}
    144 
    145 -- | @log(exp a + exp b)@, computed without intermediate overflow.
    146 --   Used by the convex-hedge two-sided combinations to update the
    147 --   running @log(K^+ + K^-)@ statistic from the two per-direction
    148 --   log-wealths.
    149 log_sum_exp :: Double -> Double -> Double
    150 log_sum_exp !a !b
    151   | a >= b    = a + log1p (exp (b - a))
    152   | otherwise = b + log1p (exp (a - b))
    153 {-# INLINE log_sum_exp #-}
    154 
    155 -- | @log 2@ as a shared constant. Used both as the initial value of
    156 --   the two-sided running sup-log-sum (since @K^+_0 + K^-_0 = 2@) and
    157 --   as the tight upper-bound slack in the fast-path skip inside
    158 --   'Numeric.Eproc.Bounded.update' /
    159 --   'Numeric.Eproc.Bernoulli.TwoSided.update'.
    160 log2_dbl :: Double
    161 log2_dbl = log 2
    162 {-# INLINE log2_dbl #-}
    163 
    164 -- | Per-bettor state. One constructor per 'Bettor' alternative; the
    165 --   constructor used in any given state matches the 'Bettor' chosen
    166 --   in the enclosing 'Config'.
    167 --
    168 --   Internal: exposed only so that the per-test 'State' types in
    169 --   "Numeric.Eproc.Bounded" and "Numeric.Eproc.Bernoulli" can share
    170 --   one implementation.
    171 data BetState =
    172     SFixed
    173   | SAdaptive
    174       {-# UNPACK #-} !Double  -- sum of z (centred observation)
    175       {-# UNPACK #-} !Double  -- sum of z^2 (for online variance)
    176       {-# UNPACK #-} !Int     -- count
    177   | SNewton
    178       {-# UNPACK #-} !Double  -- current bet lambda
    179       {-# UNPACK #-} !Double  -- running sum of per-step squared gradients
    180 
    181 -- | Per-bettor initial state.
    182 init_bet :: Bettor -> BetState
    183 init_bet b = case b of
    184   Fixed _  -> SFixed
    185   Adaptive -> SAdaptive 0 0 0
    186   Newton   -> SNewton 0 1.0e-6  -- small acc seed avoids div-by-zero
    187 {-# INLINE init_bet #-}
    188 
    189 -- | WSR (2024) Algorithm 2 ONS learning rate, @2 \/ (2 - log 3)@.
    190 ons_lr :: Double
    191 ons_lr = 2 / (2 - log 3)
    192 {-# INLINE ons_lr #-}
    193 
    194 -- | Compute the next bet 'lambda' from the bettor and its current
    195 --   state; 'lam_max' is the direction-specific safety bound. All
    196 --   strategies clip the result to @[0, lam_max]@ so the wealth
    197 --   factor stays nonnegative.
    198 bet_lambda :: Bettor -> Double -> BetState -> Double
    199 bet_lambda b !lam_max !s = case b of
    200   Fixed lam -> max 0 (min lam_max lam)
    201   Adaptive -> case s of
    202     SAdaptive !sm !sm2 !n
    203       | n == 0    -> 0
    204       | otherwise ->
    205           let !nd  = fromIntegral n
    206               !mu  = sm / nd
    207               !mu2 = mu * mu
    208               !var = max 0 (sm2 / nd - mu2)
    209               !den = var + mu2
    210               !raw = if den == 0 then 0 else mu / den
    211           in  max 0 (min lam_max raw)
    212     _ -> 0
    213   Newton -> case s of
    214     SNewton !lam _ -> lam
    215     _              -> 0
    216 {-# INLINE bet_lambda #-}
    217 
    218 -- | Update bettor state with newly observed centred value 'z'. For
    219 --   'Adaptive' this is just accumulating sums; for 'Newton' we take
    220 --   one online Newton step (with the WSR learning rate) on the
    221 --   per-step log-wealth loss @-log(1 + lambda * z)@, accumulating
    222 --   squared gradients for adaptive scaling.
    223 step_bet :: Bettor -> Double -> BetState -> Double -> BetState
    224 step_bet b !lam_max !s !z = case b of
    225   Fixed _ -> SFixed
    226   Adaptive -> case s of
    227     SAdaptive !sm !sm2 !n -> SAdaptive (sm + z) (sm2 + z * z) (n + 1)
    228     _                     -> SAdaptive z (z * z) 1
    229   Newton -> case s of
    230     SNewton !lam !acc ->
    231       let !denom = 1 + lam * z
    232           !g     = if denom == 0 then 0 else negate z / denom
    233           !acc'  = acc + g * g
    234           !lam'  = lam - ons_lr * g / acc'
    235           !clp   = max 0 (min lam_max lam')
    236       in  SNewton clp acc'
    237     _ -> SNewton 0 1.0e-6
    238 {-# INLINE step_bet #-}