ppad-eproc-0.4.0: lib/Numeric/Eproc/ConfSeq.hs
{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
-- |
-- Module: Numeric.Eproc.ConfSeq
-- Copyright: (c) 2026 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Anytime-valid confidence sequence for the mean of bounded
-- observations.
--
-- For samples @x_t@ in @[lo, hi]@ with common conditional mean
--
-- @mu = E[x_t | F_{t-1}] for all t@
--
-- (@F_{t-1}@ being the filtration generated by everything observed
-- strictly before time @t@; for i.i.d. samples this is just
-- @E[x]@), the running state yields a confidence interval @C_t@
-- after every observation, with time-uniform coverage:
--
-- @P(for all t, mu in C_t) >= 1 - alpha@
--
-- whenever @C_t@ is reported at all (see 'interval' for the empty
-- case). The guarantee holds uniformly over time, so the user may
-- inspect the interval after every observation and stop at any
-- data-dependent time -- optional stopping does not erode coverage.
--
-- The construction is the /hedged capital/ confidence sequence of
-- Waudby-Smith & Ramdas (2024), Theorem 3, evaluated over a finite
-- grid of candidate means. All arithmetic is carried out in
-- @[0, 1]@ coordinates internally; observations are mapped affinely
-- at the boundary. Each candidate @m@ runs a pair of betting
-- processes: a /positive-direction/ capital @K^+_t(m)@ wagering
-- that the mean exceeds @m@, and a /negative-direction/ capital
-- @K^-_t(m)@ wagering the reverse. The base bet is a single
-- predictable plug-in (their eq. (26)), computed once per update
-- from the running regularized mean and variance of the data and
-- shared by every candidate: it never depends on @m@, and only a
-- final truncation to @c \/ m@ (respectively @c \/ (1 - m)@), with
-- @c = 1\/2@, is candidate-specific. This @m@-freeness is what
-- makes the survivor set provably an interval (Theorem 3);
-- @m@-dependent bets can produce non-interval survivor sets (their
-- Section E.4), which is why this module does not use the library's
-- 'Numeric.Eproc.Common.Bettor' strategies.
--
-- A candidate @m@ is rejected once the max-hedge (@theta = 1\/2@)
-- capital @max(K^+_t(m), K^-_t(m)) \/ 2@ crosses @1 \/ alpha@.
-- Under the truth @m = mu@ each capital process is a nonnegative
-- supermartingale, the max is dominated by the convex combination
-- @(K^+ + K^-) \/ 2@, and Ville's inequality bounds the probability
-- that the truth is ever rejected by @alpha@. No multiplicity
-- correction across grid candidates is needed: coverage concerns
-- only the true mean's own test, and rejection of other candidates
-- merely tightens the interval.
--
-- Grid resolution is an accuracy\/cost knob. Interval endpoints are
-- quantized to the grid -- a @g@-point grid resolves them to within
-- @(hi - lo) \/ (g + 1)@ -- and per-update cost is @O(live
-- candidates)@, shrinking as evidence accumulates and candidates
-- are rejected.
--
-- == Example
--
-- Estimate the mean of a stream in @[0, 1]@ with empirical mean
-- @0.8@, at level @alpha = 0.05@ on a 100-point grid:
--
-- >>> let Right cfg = config 0.0 1.0 0.05 100
-- >>> let xs = concat (replicate 50 [1, 1, 0, 1, 1, 0, 1, 1, 1, 1])
-- >>> interval cfg (foldl' (update cfg) (initial cfg) xs)
-- Just (0.7326732673267327,0.8514851485148515)
module Numeric.Eproc.ConfSeq (
-- * Confidence-sequence configuration and state
Config
, State
, ConfigError(..)
-- * Construction
, config
, initial
-- * Streaming
, update
-- * Inspection
, interval
, samples
) where
import GHC.Float (log1p)
import Numeric.Eproc.Common (ConfigError(..), finite)
-- types ----------------------------------------------------------------------
-- | Confidence-sequence configuration. Build with 'config'.
--
-- Carries the sample bounds, the significance level, the grid
-- size, and the precomputed per-candidate rejection threshold
-- @log(2 \/ alpha)@ along with the bet numerator
-- @2 log(2 \/ alpha)@.
data Config = Config {
cfg_lo :: {-# UNPACK #-} !Double -- ^ sample lower bound
, cfg_hi :: {-# UNPACK #-} !Double -- ^ sample upper bound
, cfg_alpha :: {-# UNPACK #-} !Double -- ^ significance level
, cfg_grid :: {-# UNPACK #-} !Int -- ^ grid size @g@
, cfg_log_thresh :: {-# UNPACK #-} !Double -- ^ @log(2 \/ alpha)@
, cfg_bet_num :: {-# UNPACK #-} !Double -- ^ @2 log(2 \/ alpha)@
}
-- | One live grid candidate: its grid index and the running
-- log-capitals of the positive- and negative-direction bets.
data Point = Point
{-# UNPACK #-} !Int -- grid index j
{-# UNPACK #-} !Double -- log K^+
{-# UNPACK #-} !Double -- log K^-
-- | Streaming confidence-sequence state. Construct with 'initial'
-- and fold observations through 'update'.
--
-- Carries the sample count, the shared plug-in bettor statistics
-- (regularized running sums in @[0, 1]@ coordinates), and the
-- live grid candidates. Rejected candidates are dropped
-- permanently, so the reported intervals are nested.
--
-- Invariant: 'initial' and 'update' construct the live list fully
-- forced -- no thunks in the spine or the elements -- so a 'State'
-- in WHNF is already in normal form.
data State = State {
st_n :: {-# UNPACK #-} !Int -- ^ sample count
, st_sum_y :: {-# UNPACK #-} !Double -- ^ @sum y_i@
, st_sum_dev2 :: {-# UNPACK #-} !Double -- ^ @sum (y_i - mu_i)^2@
, st_live :: ![Point] -- ^ live grid candidates
}
-- | WSR (2024) truncation level @c = 1\/2@. Bets are capped at
-- @c \/ m@ (positive direction) and @c \/ (1 - m)@ (negative
-- direction), keeping every capital factor at least @1 - c > 0@.
trunc_c :: Double
trunc_c = 0.5
{-# INLINE trunc_c #-}
-- construction ---------------------------------------------------------------
-- | Build a 'Config' for the confidence sequence.
--
-- The candidate means form the interior grid
--
-- @m_j = lo + (j \/ (g + 1)) * (hi - lo), j = 1 .. g@
--
-- (endpoints excluded, so that in @[0, 1]@ coordinates the bet
-- truncations @c \/ m@ and @c \/ (1 - m)@ stay finite). The
-- per-candidate rejection threshold @log(2 \/ alpha)@ and the bet
-- numerator @2 log(2 \/ alpha)@ are precomputed.
--
-- Returns 'Left' with a 'ConfigError' on inputs that would leave
-- the mathematical regime: @alpha@ non-finite or outside
-- @(0, 1)@; @lo@ or @hi@ non-finite, or @lo >= hi@; or a grid
-- size below @1@.
--
-- >>> let Right cfg = config 0.0 1.0 0.05 100
config
:: Double -- ^ sample lower bound @lo@
-> Double -- ^ sample upper bound @hi@
-> Double -- ^ significance level @alpha@
-> Int -- ^ grid size @g@
-> Either ConfigError Config
config !lo !hi !alpha !g
| not (finite alpha && alpha > 0 && alpha < 1) =
Left (InvalidAlpha alpha)
| not (finite lo && finite hi && lo < hi) =
Left (InvalidBounds lo hi)
| g < 1 =
Left (InvalidGridSize g)
| otherwise = Right Config {
cfg_lo = lo
, cfg_hi = hi
, cfg_alpha = alpha
, cfg_grid = g
, cfg_log_thresh = log (2 / alpha)
, cfg_bet_num = 2 * log (2 / alpha)
}
{-# INLINE config #-}
-- | The initial 'State' for a fresh confidence sequence.
--
-- Every grid candidate starts live with both log-capitals at @0@
-- (i.e., @K^+ = K^- = 1@); the shared bettor statistics start
-- from their regularized priors (@mu_0 = 1\/2@,
-- @sigma^2_0 = 1\/4@ in @[0, 1]@ coordinates).
--
-- >>> let s0 = initial cfg
initial :: Config -> State
initial Config{..} = State {
st_n = 0
, st_sum_y = 0
, st_sum_dev2 = 0
, st_live = points 1
}
where
-- built eagerly: the tail is forced before consing, so the
-- whole list is in normal form on construction.
points !j
| j > cfg_grid = []
| otherwise =
let !p = Point j 0 0
!rest = points (j + 1)
in p : rest
{-# INLINE initial #-}
-- streaming ------------------------------------------------------------------
-- | Fold one observation into the running 'State'.
--
-- Maps the observation to @[0, 1]@ coordinates via
-- @y = (x - lo) \/ (hi - lo)@ and computes the shared predictable
-- plug-in bet from the statistics accumulated through the
-- /previous/ step (Waudby-Smith & Ramdas (2024), eq. (26)):
--
-- @lambda_t = min c (sqrt (2 log(2 \/ alpha)
-- \/ (sigma^2_{t-1} * t * log(1 + t))))@
--
-- with @c = 1\/2@. The bet is computed once and shared across all
-- live candidates -- its independence from @m@ is what keeps the
-- survivor set an interval. Each live candidate @m@ then updates
-- its pair of log-capitals with the truncated bets
-- @min lambda_t (c \/ m)@ and @min lambda_t (c \/ (1 - m))@, and
-- is dropped iff @max(log K^+, log K^-)@ has reached
-- @log(2 \/ alpha)@. Finally @y@ is folded into the shared
-- statistics, preserving predictability of the next bet.
--
-- /Precondition/: @x@ must lie in the @[lo, hi]@ interval given
-- to 'config'. The coverage guarantee of the sequence depends on
-- it. Out-of-range observations can drive a capital factor
-- negative, taking the construction out of the supermartingale
-- regime entirely; the function does not check for this.
--
-- >>> let s1 = update cfg s0 0.7
update :: Config -> State -> Double -> State
update Config{..} State{..} !x =
let !y = (x - cfg_lo) / (cfg_hi - cfg_lo)
!t = st_n + 1
!td = fromIntegral t
!gp1 = fromIntegral (cfg_grid + 1)
-- sigma^2_{t-1} = (1/4 + sum_{i<=t-1} (y_i - mu_i)^2) / t
!sig2 = (0.25 + st_sum_dev2) / td
!lam = min trunc_c
(sqrt (cfg_bet_num / (sig2 * td * log1p td)))
-- built eagerly, as in 'initial': the tail is forced before
-- consing, so the new live list is in normal form on
-- construction.
go [] = []
go (Point j lp ln : ps) =
let !m = fromIntegral j / gp1
!d = y - m
!lp' = lp + log1p (min lam (trunc_c / m) * d)
!ln' = ln + log1p (negate (min lam (trunc_c / (1 - m)))
* d)
!rest = go ps
in if max lp' ln' >= cfg_log_thresh
then rest
else Point j lp' ln' : rest
!live = go st_live
-- fold y into the shared statistics only now: the bet above
-- used statistics through t-1, so predictability holds. the
-- deviation at step t uses the current-inclusive mean mu_t.
!sum_y' = st_sum_y + y
!mu = (0.5 + sum_y') / (td + 1)
!dev = y - mu
!dev2' = st_sum_dev2 + dev * dev
in State t sum_y' dev2' live
{-# INLINE update #-}
-- inspection -----------------------------------------------------------------
-- | The current confidence interval, in the original @[lo, hi]@
-- coordinates.
--
-- The interval spans the surviving grid candidates, widened by
-- one grid step at each end (or clamped to @lo@ \/ @hi@ at the
-- grid's edges). The widening is what makes off-grid true means
-- safe: Theorem 3 guarantees the ideal continuum survivor set is
-- an interval, so its endpoints are bracketed by the nearest
-- /rejected/ grid candidates, and reporting those sentinels
-- yields a superset of the continuum interval. Whenever the
-- result is 'Just', it therefore covers the true mean uniformly
-- over time with probability at least @1 - alpha@ -- no
-- multiplicity correction across candidates is needed, since
-- coverage concerns only the true mean's own test.
--
-- 'Nothing' means every grid candidate has been rejected: the
-- evidence has resolved the mean below the grid's resolution.
-- For a true mean lying exactly on the grid this has probability
-- at most @alpha@ (its own test must have rejected). For an
-- off-grid true mean it additionally occurs once the continuum
-- survivor interval shrinks inside a single grid cell -- a
-- quantization horizon far beyond the point where the reported
-- width is comparable to the grid spacing. Treat 'Nothing' as a
-- signal to rerun with a larger grid, not as an inference.
--
-- >>> interval cfg (initial cfg)
-- Just (0.0,1.0)
interval :: Config -> State -> Maybe (Double, Double)
interval Config{..} State{..} = case st_live of
[] -> Nothing
(Point j0 _ _ : ps) ->
let !jmin = foldl' (\acc (Point j _ _) -> min acc j) j0 ps
!jmax = foldl' (\acc (Point j _ _) -> max acc j) j0 ps
!gp1 = fromIntegral (cfg_grid + 1)
!w = cfg_hi - cfg_lo
!l | jmin == 1 = cfg_lo
| otherwise =
cfg_lo + fromIntegral (jmin - 1) / gp1 * w
!u | jmax == cfg_grid = cfg_hi
| otherwise =
cfg_lo + fromIntegral (jmax + 1) / gp1 * w
in Just (l, u)
{-# INLINE interval #-}
-- | The number of samples consumed so far.
--
-- >>> samples s0
-- 0
samples :: State -> Int
samples = st_n
{-# INLINE samples #-}