ConfSeq.hs (13079B)
1 {-# OPTIONS_HADDOCK prune #-} 2 {-# LANGUAGE BangPatterns #-} 3 {-# LANGUAGE RecordWildCards #-} 4 5 -- | 6 -- Module: Numeric.Eproc.ConfSeq 7 -- Copyright: (c) 2026 Jared Tobin 8 -- License: MIT 9 -- Maintainer: Jared Tobin <jared@ppad.tech> 10 -- 11 -- Anytime-valid confidence sequence for the mean of bounded 12 -- observations. 13 -- 14 -- For samples @x_t@ in @[lo, hi]@ with common conditional mean 15 -- 16 -- @mu = E[x_t | F_{t-1}] for all t@ 17 -- 18 -- (@F_{t-1}@ being the filtration generated by everything observed 19 -- strictly before time @t@; for i.i.d. samples this is just 20 -- @E[x]@), the running state yields a confidence interval @C_t@ 21 -- after every observation, with time-uniform coverage: 22 -- 23 -- @P(for all t, mu in C_t) >= 1 - alpha@ 24 -- 25 -- whenever @C_t@ is reported at all (see 'interval' for the empty 26 -- case). The guarantee holds uniformly over time, so the user may 27 -- inspect the interval after every observation and stop at any 28 -- data-dependent time -- optional stopping does not erode coverage. 29 -- 30 -- The construction is the /hedged capital/ confidence sequence of 31 -- Waudby-Smith & Ramdas (2024), Theorem 3, evaluated over a finite 32 -- grid of candidate means. All arithmetic is carried out in 33 -- @[0, 1]@ coordinates internally; observations are mapped affinely 34 -- at the boundary. Each candidate @m@ runs a pair of betting 35 -- processes: a /positive-direction/ capital @K^+_t(m)@ wagering 36 -- that the mean exceeds @m@, and a /negative-direction/ capital 37 -- @K^-_t(m)@ wagering the reverse. The base bet is a single 38 -- predictable plug-in (their eq. (26)), computed once per update 39 -- from the running regularized mean and variance of the data and 40 -- shared by every candidate: it never depends on @m@, and only a 41 -- final truncation to @c \/ m@ (respectively @c \/ (1 - m)@), with 42 -- @c = 1\/2@, is candidate-specific. This @m@-freeness is what 43 -- makes the survivor set provably an interval (Theorem 3); 44 -- @m@-dependent bets can produce non-interval survivor sets (their 45 -- Section E.4), which is why this module does not use the library's 46 -- 'Numeric.Eproc.Common.Bettor' strategies. 47 -- 48 -- A candidate @m@ is rejected once the max-hedge (@theta = 1\/2@) 49 -- capital @max(K^+_t(m), K^-_t(m)) \/ 2@ crosses @1 \/ alpha@. 50 -- Under the truth @m = mu@ each capital process is a nonnegative 51 -- supermartingale, the max is dominated by the convex combination 52 -- @(K^+ + K^-) \/ 2@, and Ville's inequality bounds the probability 53 -- that the truth is ever rejected by @alpha@. No multiplicity 54 -- correction across grid candidates is needed: coverage concerns 55 -- only the true mean's own test, and rejection of other candidates 56 -- merely tightens the interval. 57 -- 58 -- Grid resolution is an accuracy\/cost knob. Interval endpoints are 59 -- quantized to the grid -- a @g@-point grid resolves them to within 60 -- @(hi - lo) \/ (g + 1)@ -- and per-update cost is @O(live 61 -- candidates)@, shrinking as evidence accumulates and candidates 62 -- are rejected. 63 -- 64 -- == Example 65 -- 66 -- Estimate the mean of a stream in @[0, 1]@ with empirical mean 67 -- @0.8@, at level @alpha = 0.05@ on a 100-point grid: 68 -- 69 -- >>> let Right cfg = config 0.0 1.0 0.05 100 70 -- >>> let xs = concat (replicate 50 [1, 1, 0, 1, 1, 0, 1, 1, 1, 1]) 71 -- >>> interval cfg (foldl' (update cfg) (initial cfg) xs) 72 -- Just (0.7326732673267327,0.8613861386138614) 73 74 module Numeric.Eproc.ConfSeq ( 75 -- * Confidence-sequence configuration and state 76 Config 77 , State 78 , ConfigError(..) 79 80 -- * Construction 81 , config 82 , initial 83 84 -- * Streaming 85 , update 86 87 -- * Inspection 88 , interval 89 , samples 90 ) where 91 92 import GHC.Float (log1p) 93 import Numeric.Eproc.Common (ConfigError(..), finite) 94 95 -- types ---------------------------------------------------------------------- 96 97 -- | Confidence-sequence configuration. Build with 'config'. 98 -- 99 -- Carries the sample bounds, the significance level, the grid 100 -- size, and the precomputed per-candidate rejection threshold 101 -- @log(2 \/ alpha)@ along with the bet numerator 102 -- @2 log(2 \/ alpha)@. 103 data Config = Config { 104 cfg_lo :: {-# UNPACK #-} !Double -- ^ sample lower bound 105 , cfg_hi :: {-# UNPACK #-} !Double -- ^ sample upper bound 106 , cfg_alpha :: {-# UNPACK #-} !Double -- ^ significance level 107 , cfg_grid :: {-# UNPACK #-} !Int -- ^ grid size @g@ 108 , cfg_log_thresh :: {-# UNPACK #-} !Double -- ^ @log(2 \/ alpha)@ 109 , cfg_bet_num :: {-# UNPACK #-} !Double -- ^ @2 log(2 \/ alpha)@ 110 } 111 112 -- | One live grid candidate: its grid index and the running 113 -- log-capitals of the positive- and negative-direction bets. 114 data Point = Point 115 {-# UNPACK #-} !Int -- grid index j 116 {-# UNPACK #-} !Double -- log K^+ 117 {-# UNPACK #-} !Double -- log K^- 118 119 -- | Streaming confidence-sequence state. Construct with 'initial' 120 -- and fold observations through 'update'. 121 -- 122 -- Carries the sample count, the shared plug-in bettor statistics 123 -- (regularized running sums in @[0, 1]@ coordinates), and the 124 -- live grid candidates. Rejected candidates are dropped 125 -- permanently, so the reported intervals are nested. 126 -- 127 -- Invariant: 'initial' and 'update' construct the live list fully 128 -- forced -- no thunks in the spine or the elements -- so a 'State' 129 -- in WHNF is already in normal form. 130 data State = State { 131 st_n :: {-# UNPACK #-} !Int -- ^ sample count 132 , st_sum_y :: {-# UNPACK #-} !Double -- ^ @sum y_i@ 133 , st_sum_dev2 :: {-# UNPACK #-} !Double -- ^ @sum (y_i - mu_i)^2@ 134 , st_live :: ![Point] -- ^ live grid candidates 135 } 136 137 -- | WSR (2024) truncation level @c = 1\/2@. Bets are capped at 138 -- @c \/ m@ (positive direction) and @c \/ (1 - m)@ (negative 139 -- direction), keeping every capital factor at least @1 - c > 0@. 140 trunc_c :: Double 141 trunc_c = 0.5 142 {-# INLINE trunc_c #-} 143 144 -- construction --------------------------------------------------------------- 145 146 -- | Build a 'Config' for the confidence sequence. 147 -- 148 -- The candidate means form the interior grid 149 -- 150 -- @m_j = lo + (j \/ (g + 1)) * (hi - lo), j = 1 .. g@ 151 -- 152 -- (endpoints excluded, so that in @[0, 1]@ coordinates the bet 153 -- truncations @c \/ m@ and @c \/ (1 - m)@ stay finite). The 154 -- per-candidate rejection threshold @log(2 \/ alpha)@ and the bet 155 -- numerator @2 log(2 \/ alpha)@ are precomputed. 156 -- 157 -- Returns 'Left' with a 'ConfigError' on inputs that would leave 158 -- the mathematical regime: @alpha@ non-finite or outside 159 -- @(0, 1)@; @lo@ or @hi@ non-finite, or @lo >= hi@; or a grid 160 -- size below @1@. 161 -- 162 -- >>> let Right cfg = config 0.0 1.0 0.05 100 163 config 164 :: Double -- ^ sample lower bound @lo@ 165 -> Double -- ^ sample upper bound @hi@ 166 -> Double -- ^ significance level @alpha@ 167 -> Int -- ^ grid size @g@ 168 -> Either ConfigError Config 169 config !lo !hi !alpha !g 170 | not (finite alpha && alpha > 0 && alpha < 1) = 171 Left (InvalidAlpha alpha) 172 | not (finite lo && finite hi && lo < hi) = 173 Left (InvalidBounds lo hi) 174 | g < 1 = 175 Left (InvalidGridSize g) 176 | otherwise = Right Config { 177 cfg_lo = lo 178 , cfg_hi = hi 179 , cfg_alpha = alpha 180 , cfg_grid = g 181 , cfg_log_thresh = log (2 / alpha) 182 , cfg_bet_num = 2 * log (2 / alpha) 183 } 184 {-# INLINE config #-} 185 186 -- | The initial 'State' for a fresh confidence sequence. 187 -- 188 -- Every grid candidate starts live with both log-capitals at @0@ 189 -- (i.e., @K^+ = K^- = 1@); the shared bettor statistics start 190 -- from their regularized priors (@mu_0 = 1\/2@, 191 -- @sigma^2_0 = 1\/4@ in @[0, 1]@ coordinates). 192 -- 193 -- >>> let s0 = initial cfg 194 initial :: Config -> State 195 initial Config{..} = State { 196 st_n = 0 197 , st_sum_y = 0 198 , st_sum_dev2 = 0 199 , st_live = points 1 200 } 201 where 202 -- built eagerly: the tail is forced before consing, so the 203 -- whole list is in normal form on construction. 204 points !j 205 | j > cfg_grid = [] 206 | otherwise = 207 let !p = Point j 0 0 208 !rest = points (j + 1) 209 in p : rest 210 {-# INLINE initial #-} 211 212 -- streaming ------------------------------------------------------------------ 213 214 -- | Fold one observation into the running 'State'. 215 -- 216 -- Maps the observation to @[0, 1]@ coordinates via 217 -- @y = (x - lo) \/ (hi - lo)@ and computes the shared predictable 218 -- plug-in bet from the statistics accumulated through the 219 -- /previous/ step (Waudby-Smith & Ramdas (2024), eq. (26)): 220 -- 221 -- @lambda_t = sqrt (2 log(2 \/ alpha) 222 -- \/ (sigma^2_{t-1} * t * log(1 + t)))@ 223 -- 224 -- The bet is computed once and shared across all 225 -- live candidates -- its independence from @m@ is what keeps the 226 -- survivor set an interval. Each live candidate @m@ then updates 227 -- its pair of log-capitals with the truncated bets 228 -- @min lambda_t (c \/ m)@ and @min lambda_t (c \/ (1 - m))@, 229 -- with @c = 1\/2@, and 230 -- is dropped iff @max(log K^+, log K^-)@ has reached 231 -- @log(2 \/ alpha)@. Finally @y@ is folded into the shared 232 -- statistics, preserving predictability of the next bet. 233 -- 234 -- /Precondition/: @x@ must lie in the @[lo, hi]@ interval given 235 -- to 'config'. The coverage guarantee of the sequence depends on 236 -- it. Out-of-range observations can drive a capital factor 237 -- negative, taking the construction out of the supermartingale 238 -- regime entirely; the function does not check for this. 239 -- 240 -- >>> let s1 = update cfg s0 0.7 241 update :: Config -> State -> Double -> State 242 update Config{..} State{..} !x = 243 let !y = (x - cfg_lo) / (cfg_hi - cfg_lo) 244 !t = st_n + 1 245 !td = fromIntegral t 246 !gp1 = fromIntegral (cfg_grid + 1) 247 -- sigma^2_{t-1} = (1/4 + sum_{i<=t-1} (y_i - mu_i)^2) / t 248 !sig2 = (0.25 + st_sum_dev2) / td 249 !lam = sqrt (cfg_bet_num / (sig2 * td * log1p td)) 250 -- built eagerly, as in 'initial': the tail is forced before 251 -- consing, so the new live list is in normal form on 252 -- construction. 253 go [] = [] 254 go (Point j lp ln : ps) = 255 let !m = fromIntegral j / gp1 256 !d = y - m 257 !lp' = lp + log1p (min lam (trunc_c / m) * d) 258 !ln' = ln + log1p (negate (min lam (trunc_c / (1 - m))) 259 * d) 260 !rest = go ps 261 in if max lp' ln' >= cfg_log_thresh 262 then rest 263 else Point j lp' ln' : rest 264 !live = go st_live 265 -- fold y into the shared statistics only now: the bet above 266 -- used statistics through t-1, so predictability holds. the 267 -- deviation at step t uses the current-inclusive mean mu_t. 268 !sum_y' = st_sum_y + y 269 !mu = (0.5 + sum_y') / (td + 1) 270 !dev = y - mu 271 !dev2' = st_sum_dev2 + dev * dev 272 in State t sum_y' dev2' live 273 {-# INLINE update #-} 274 275 -- inspection ----------------------------------------------------------------- 276 277 -- | The current confidence interval, in the original @[lo, hi]@ 278 -- coordinates. 279 -- 280 -- The interval spans the surviving grid candidates, widened by 281 -- one grid step at each end (or clamped to @lo@ \/ @hi@ at the 282 -- grid's edges). The widening is what makes off-grid true means 283 -- safe: Theorem 3 guarantees the ideal continuum survivor set is 284 -- an interval, so its endpoints are bracketed by the nearest 285 -- /rejected/ grid candidates, and reporting those sentinels 286 -- yields a superset of the continuum interval. Whenever the 287 -- result is 'Just', it therefore covers the true mean uniformly 288 -- over time with probability at least @1 - alpha@ -- no 289 -- multiplicity correction across candidates is needed, since 290 -- coverage concerns only the true mean's own test. 291 -- 292 -- 'Nothing' means every grid candidate has been rejected: the 293 -- evidence has resolved the mean below the grid's resolution. 294 -- For a true mean lying exactly on the grid this has probability 295 -- at most @alpha@ (its own test must have rejected). For an 296 -- off-grid true mean it additionally occurs once the continuum 297 -- survivor interval shrinks inside a single grid cell -- a 298 -- quantization horizon far beyond the point where the reported 299 -- width is comparable to the grid spacing. Treat 'Nothing' as a 300 -- signal to rerun with a larger grid, not as an inference. 301 -- 302 -- >>> interval cfg (initial cfg) 303 -- Just (0.0,1.0) 304 interval :: Config -> State -> Maybe (Double, Double) 305 interval Config{..} State{..} = case st_live of 306 [] -> Nothing 307 (Point j0 _ _ : ps) -> 308 let !jmin = foldl' (\acc (Point j _ _) -> min acc j) j0 ps 309 !jmax = foldl' (\acc (Point j _ _) -> max acc j) j0 ps 310 !gp1 = fromIntegral (cfg_grid + 1) 311 !w = cfg_hi - cfg_lo 312 !l | jmin == 1 = cfg_lo 313 | otherwise = 314 cfg_lo + fromIntegral (jmin - 1) / gp1 * w 315 !u | jmax == cfg_grid = cfg_hi 316 | otherwise = 317 cfg_lo + fromIntegral (jmax + 1) / gp1 * w 318 in Just (l, u) 319 {-# INLINE interval #-} 320 321 -- | The number of samples consumed so far. 322 -- 323 -- >>> samples s0 324 -- 0 325 samples :: State -> Int 326 samples = st_n 327 {-# INLINE samples #-}