ConfSeq.hs (13115B)
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 qualified Data.List as L 93 import GHC.Float (log1p) 94 import Numeric.Eproc.Common (ConfigError(..), finite) 95 96 -- types ---------------------------------------------------------------------- 97 98 -- | Confidence-sequence configuration. Build with 'config'. 99 -- 100 -- Carries the sample bounds, the significance level, the grid 101 -- size, and the precomputed per-candidate rejection threshold 102 -- @log(2 \/ alpha)@ along with the bet numerator 103 -- @2 log(2 \/ alpha)@. 104 data Config = Config { 105 cfg_lo :: {-# UNPACK #-} !Double -- ^ sample lower bound 106 , cfg_hi :: {-# UNPACK #-} !Double -- ^ sample upper bound 107 , cfg_alpha :: {-# UNPACK #-} !Double -- ^ significance level 108 , cfg_grid :: {-# UNPACK #-} !Int -- ^ grid size @g@ 109 , cfg_log_thresh :: {-# UNPACK #-} !Double -- ^ @log(2 \/ alpha)@ 110 , cfg_bet_num :: {-# UNPACK #-} !Double -- ^ @2 log(2 \/ alpha)@ 111 } 112 113 -- | One live grid candidate: its grid index and the running 114 -- log-capitals of the positive- and negative-direction bets. 115 data Point = Point 116 {-# UNPACK #-} !Int -- grid index j 117 {-# UNPACK #-} !Double -- log K^+ 118 {-# UNPACK #-} !Double -- log K^- 119 120 -- | Streaming confidence-sequence state. Construct with 'initial' 121 -- and fold observations through 'update'. 122 -- 123 -- Carries the sample count, the shared plug-in bettor statistics 124 -- (regularized running sums in @[0, 1]@ coordinates), and the 125 -- live grid candidates. Rejected candidates are dropped 126 -- permanently, so the reported intervals are nested. 127 -- 128 -- Invariant: 'initial' and 'update' construct the live list fully 129 -- forced -- no thunks in the spine or the elements -- so a 'State' 130 -- in WHNF is already in normal form. 131 data State = State { 132 st_n :: {-# UNPACK #-} !Int -- ^ sample count 133 , st_sum_y :: {-# UNPACK #-} !Double -- ^ @sum y_i@ 134 , st_sum_dev2 :: {-# UNPACK #-} !Double -- ^ @sum (y_i - mu_i)^2@ 135 , st_live :: ![Point] -- ^ live grid candidates 136 } 137 138 -- | WSR (2024) truncation level @c = 1\/2@. Bets are capped at 139 -- @c \/ m@ (positive direction) and @c \/ (1 - m)@ (negative 140 -- direction), keeping every capital factor at least @1 - c > 0@. 141 trunc_c :: Double 142 trunc_c = 0.5 143 {-# INLINE trunc_c #-} 144 145 -- construction --------------------------------------------------------------- 146 147 -- | Build a 'Config' for the confidence sequence. 148 -- 149 -- The candidate means form the interior grid 150 -- 151 -- @m_j = lo + (j \/ (g + 1)) * (hi - lo), j = 1 .. g@ 152 -- 153 -- (endpoints excluded, so that in @[0, 1]@ coordinates the bet 154 -- truncations @c \/ m@ and @c \/ (1 - m)@ stay finite). The 155 -- per-candidate rejection threshold @log(2 \/ alpha)@ and the bet 156 -- numerator @2 log(2 \/ alpha)@ are precomputed. 157 -- 158 -- Returns 'Left' with a 'ConfigError' on inputs that would leave 159 -- the mathematical regime: @alpha@ non-finite or outside 160 -- @(0, 1)@; @lo@ or @hi@ non-finite, or @lo >= hi@; or a grid 161 -- size below @1@. 162 -- 163 -- >>> let Right cfg = config 0.0 1.0 0.05 100 164 config 165 :: Double -- ^ sample lower bound @lo@ 166 -> Double -- ^ sample upper bound @hi@ 167 -> Double -- ^ significance level @alpha@ 168 -> Int -- ^ grid size @g@ 169 -> Either ConfigError Config 170 config !lo !hi !alpha !g 171 | not (finite alpha && alpha > 0 && alpha < 1) = 172 Left (InvalidAlpha alpha) 173 | not (finite lo && finite hi && lo < hi) = 174 Left (InvalidBounds lo hi) 175 | g < 1 = 176 Left (InvalidGridSize g) 177 | otherwise = Right Config { 178 cfg_lo = lo 179 , cfg_hi = hi 180 , cfg_alpha = alpha 181 , cfg_grid = g 182 , cfg_log_thresh = log (2 / alpha) 183 , cfg_bet_num = 2 * log (2 / alpha) 184 } 185 {-# INLINE config #-} 186 187 -- | The initial 'State' for a fresh confidence sequence. 188 -- 189 -- Every grid candidate starts live with both log-capitals at @0@ 190 -- (i.e., @K^+ = K^- = 1@); the shared bettor statistics start 191 -- from their regularized priors (@mu_0 = 1\/2@, 192 -- @sigma^2_0 = 1\/4@ in @[0, 1]@ coordinates). 193 -- 194 -- >>> let s0 = initial cfg 195 initial :: Config -> State 196 initial Config{..} = State { 197 st_n = 0 198 , st_sum_y = 0 199 , st_sum_dev2 = 0 200 , st_live = points 1 201 } 202 where 203 -- built eagerly: the tail is forced before consing, so the 204 -- whole list is in normal form on construction. 205 points !j 206 | j > cfg_grid = [] 207 | otherwise = 208 let !p = Point j 0 0 209 !rest = points (j + 1) 210 in p : rest 211 {-# INLINE initial #-} 212 213 -- streaming ------------------------------------------------------------------ 214 215 -- | Fold one observation into the running 'State'. 216 -- 217 -- Maps the observation to @[0, 1]@ coordinates via 218 -- @y = (x - lo) \/ (hi - lo)@ and computes the shared predictable 219 -- plug-in bet from the statistics accumulated through the 220 -- /previous/ step (Waudby-Smith & Ramdas (2024), eq. (26)): 221 -- 222 -- @lambda_t = sqrt (2 log(2 \/ alpha) 223 -- \/ (sigma^2_{t-1} * t * log(1 + t)))@ 224 -- 225 -- The bet is computed once and shared across all 226 -- live candidates -- its independence from @m@ is what keeps the 227 -- survivor set an interval. Each live candidate @m@ then updates 228 -- its pair of log-capitals with the truncated bets 229 -- @min lambda_t (c \/ m)@ and @min lambda_t (c \/ (1 - m))@, 230 -- with @c = 1\/2@, and 231 -- is dropped iff @max(log K^+, log K^-)@ has reached 232 -- @log(2 \/ alpha)@. Finally @y@ is folded into the shared 233 -- statistics, preserving predictability of the next bet. 234 -- 235 -- /Precondition/: @x@ must lie in the @[lo, hi]@ interval given 236 -- to 'config'. The coverage guarantee of the sequence depends on 237 -- it. Out-of-range observations can drive a capital factor 238 -- negative, taking the construction out of the supermartingale 239 -- regime entirely; the function does not check for this. 240 -- 241 -- >>> let s1 = update cfg s0 0.7 242 update :: Config -> State -> Double -> State 243 update Config{..} State{..} !x = 244 let !y = (x - cfg_lo) / (cfg_hi - cfg_lo) 245 !t = st_n + 1 246 !td = fromIntegral t 247 !gp1 = fromIntegral (cfg_grid + 1) 248 -- sigma^2_{t-1} = (1/4 + sum_{i<=t-1} (y_i - mu_i)^2) / t 249 !sig2 = (0.25 + st_sum_dev2) / td 250 !lam = sqrt (cfg_bet_num / (sig2 * td * log1p td)) 251 -- built eagerly, as in 'initial': the tail is forced before 252 -- consing, so the new live list is in normal form on 253 -- construction. 254 go [] = [] 255 go (Point j lp ln : ps) = 256 let !m = fromIntegral j / gp1 257 !d = y - m 258 !lp' = lp + log1p (min lam (trunc_c / m) * d) 259 !ln' = ln + log1p (negate (min lam (trunc_c / (1 - m))) 260 * d) 261 !rest = go ps 262 in if max lp' ln' >= cfg_log_thresh 263 then rest 264 else Point j lp' ln' : rest 265 !live = go st_live 266 -- fold y into the shared statistics only now: the bet above 267 -- used statistics through t-1, so predictability holds. the 268 -- deviation at step t uses the current-inclusive mean mu_t. 269 !sum_y' = st_sum_y + y 270 !mu = (0.5 + sum_y') / (td + 1) 271 !dev = y - mu 272 !dev2' = st_sum_dev2 + dev * dev 273 in State t sum_y' dev2' live 274 {-# INLINE update #-} 275 276 -- inspection ----------------------------------------------------------------- 277 278 -- | The current confidence interval, in the original @[lo, hi]@ 279 -- coordinates. 280 -- 281 -- The interval spans the surviving grid candidates, widened by 282 -- one grid step at each end (or clamped to @lo@ \/ @hi@ at the 283 -- grid's edges). The widening is what makes off-grid true means 284 -- safe: Theorem 3 guarantees the ideal continuum survivor set is 285 -- an interval, so its endpoints are bracketed by the nearest 286 -- /rejected/ grid candidates, and reporting those sentinels 287 -- yields a superset of the continuum interval. Whenever the 288 -- result is 'Just', it therefore covers the true mean uniformly 289 -- over time with probability at least @1 - alpha@ -- no 290 -- multiplicity correction across candidates is needed, since 291 -- coverage concerns only the true mean's own test. 292 -- 293 -- 'Nothing' means every grid candidate has been rejected: the 294 -- evidence has resolved the mean below the grid's resolution. 295 -- For a true mean lying exactly on the grid this has probability 296 -- at most @alpha@ (its own test must have rejected). For an 297 -- off-grid true mean it additionally occurs once the continuum 298 -- survivor interval shrinks inside a single grid cell -- a 299 -- quantization horizon far beyond the point where the reported 300 -- width is comparable to the grid spacing. Treat 'Nothing' as a 301 -- signal to rerun with a larger grid, not as an inference. 302 -- 303 -- >>> interval cfg (initial cfg) 304 -- Just (0.0,1.0) 305 interval :: Config -> State -> Maybe (Double, Double) 306 interval Config{..} State{..} = case st_live of 307 [] -> Nothing 308 (Point j0 _ _ : ps) -> 309 let !jmin = L.foldl' (\acc (Point j _ _) -> min acc j) j0 ps 310 !jmax = L.foldl' (\acc (Point j _ _) -> max acc j) j0 ps 311 !gp1 = fromIntegral (cfg_grid + 1) 312 !w = cfg_hi - cfg_lo 313 !l | jmin == 1 = cfg_lo 314 | otherwise = 315 cfg_lo + fromIntegral (jmin - 1) / gp1 * w 316 !u | jmax == cfg_grid = cfg_hi 317 | otherwise = 318 cfg_lo + fromIntegral (jmax + 1) / gp1 * w 319 in Just (l, u) 320 {-# INLINE interval #-} 321 322 -- | The number of samples consumed so far. 323 -- 324 -- >>> samples s0 325 -- 0 326 samples :: State -> Int 327 samples = st_n 328 {-# INLINE samples #-}