commit a230e6558be4050da6ecb1a799d9993a9c957152
parent f1318f8b4511eae8fe5cbe9af2ab6b988c091dd2
Author: Jared Tobin <jared@jtobin.io>
Date: Tue, 2 Jun 2026 21:05:41 -0230
perf: close Bettor sum, fix per-step Rational conversion
Refactor Bettor from a record of strategy functions (init/step/bet)
parameterised over an opaque state type to a closed sum type with
three constructors (Fixed Double, Agrapa, Ons). Mean.update now
case-dispatches on the bettor; with INLINE, GHC fully specialises
each strategy at the call site. Per-direction bettor state lives in a
separate closed BetState sum carried by Mean.State.
Mean.Config and Mean.State lose their phantom 's' parameter; the
re-exports in Statistics.EProcess drop AGRAPA/ONS/fixed/agrapa/ons in
favour of Bettor(..). Breaking change to the (brand-new) public API.
While here: replace 'tiny = 1.0e-300 :: Double' with the MagicHash
form 'D# 1.0e-300##'. The fractional literal compiles as
'fromRational (1.0e-300 :: Rational)' and GHC was not folding the
conversion, leaving a $wrationalToDouble call in every update step.
This was the dominant cost once the closed-sum refactor unlocked
inlining.
Combined effect (LLVM, 1000-sample fold):
Mean fold/fixed : 128.9 us -> 4.84 us (27x)
Mean fold/agrapa : 142.4 us -> 15.67 us ( 9x)
Mean fold/ons : 133.2 us -> 14.43 us ( 9x)
TS fold/fixed : 129.2 us -> 5.33 us (24x)
TS fold/agrapa : 155.7 us -> 9.92 us (16x)
TS fold/ons : 136.9 us -> 16.31 us ( 8x)
Allocation per 1000-step fold drops from ~850 KB to <33 KB; the
fixed-bettor inner loop is fully fused (48 B total).
All 11 tests pass.
Diffstat:
7 files changed, 223 insertions(+), 246 deletions(-)
diff --git a/bench/Main.hs b/bench/Main.hs
@@ -5,7 +5,6 @@ module Main where
import Control.DeepSeq
import qualified Statistics.EProcess as E
-import qualified Statistics.EProcess.Bettor as B
import qualified Statistics.EProcess.Mean as M
import qualified Statistics.EProcess.TwoSample as TS
import Criterion.Main
@@ -13,11 +12,9 @@ import Criterion.Main
-- all relevant fields are strict (and UNPACK'd for the doubles), so
-- WHNF == NF for these types. orphan instances keep the library API
-- untouched.
-instance NFData B.AGRAPA where rnf !_ = ()
-instance NFData B.ONS where rnf !_ = ()
-instance NFData (M.State s) where rnf !_ = ()
-instance NFData (TS.State s) where rnf !_ = ()
-instance NFData M.Verdict where rnf !_ = ()
+instance NFData M.State where rnf !_ = ()
+instance NFData TS.State where rnf !_ = ()
+instance NFData M.Verdict where rnf !_ = ()
main :: IO ()
main = defaultMain [
@@ -29,12 +26,9 @@ main = defaultMain [
update :: Benchmark
update =
- let !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (const (E.fixed 0.5))
- :: M.Config ()
- !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.agrapa
- :: M.Config B.AGRAPA
- !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.ons
- :: M.Config B.ONS
+ let !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (E.Fixed 0.5)
+ !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.Agrapa
+ !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
!stF = M.initial cfgF
!stA = M.initial cfgA
!stO = M.initial cfgO
@@ -47,7 +41,7 @@ update =
decide :: Benchmark
decide =
- let !cfg = M.config 0.5 0.0 1.0 1.0e-3 E.ons :: M.Config B.ONS
+ let !cfg = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
!st = M.initial cfg
in bgroup "Mean.decide" [
bench "initial state" $ nf (M.decide cfg) st
@@ -56,12 +50,9 @@ decide =
stream :: Benchmark
stream =
let !xs = force (take 1000 (cycle [0.3, 0.7]))
- !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (const (E.fixed 0.5))
- :: M.Config ()
- !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.agrapa
- :: M.Config B.AGRAPA
- !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.ons
- :: M.Config B.ONS
+ !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (E.Fixed 0.5)
+ !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.Agrapa
+ !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
runM cfg = foldl' (M.update cfg) (M.initial cfg)
in bgroup "Mean.update (1000-sample fold)" [
bench "fixed" $ nf (runM cfgF) xs
@@ -72,12 +63,9 @@ stream =
twosample :: Benchmark
twosample =
let !ps = force (take 1000 (cycle [(0.3, 0.7), (0.7, 0.3)]))
- !cfgF = TS.config 0.0 1.0 1.0e-3 (const (E.fixed 0.5))
- :: TS.Config ()
- !cfgA = TS.config 0.0 1.0 1.0e-3 E.agrapa
- :: TS.Config B.AGRAPA
- !cfgO = TS.config 0.0 1.0 1.0e-3 E.ons
- :: TS.Config B.ONS
+ !cfgF = TS.config 0.0 1.0 1.0e-3 (E.Fixed 0.5)
+ !cfgA = TS.config 0.0 1.0 1.0e-3 E.Agrapa
+ !cfgO = TS.config 0.0 1.0 1.0e-3 E.Ons
runT cfg = foldl' (TS.update cfg) (TS.initial cfg)
in bgroup "TwoSample.update (1000-sample fold)" [
bench "fixed" $ nf (runT cfgF) ps
diff --git a/bench/Weight.hs b/bench/Weight.hs
@@ -5,16 +5,13 @@ module Main where
import Control.DeepSeq
import qualified Statistics.EProcess as E
-import qualified Statistics.EProcess.Bettor as B
import qualified Statistics.EProcess.Mean as M
import qualified Statistics.EProcess.TwoSample as TS
import Weigh
-instance NFData B.AGRAPA where rnf !_ = ()
-instance NFData B.ONS where rnf !_ = ()
-instance NFData (M.State s) where rnf !_ = ()
-instance NFData (TS.State s) where rnf !_ = ()
-instance NFData M.Verdict where rnf !_ = ()
+instance NFData M.State where rnf !_ = ()
+instance NFData TS.State where rnf !_ = ()
+instance NFData M.Verdict where rnf !_ = ()
-- note that 'weigh' doesn't work properly in a repl
main :: IO ()
@@ -26,12 +23,9 @@ main = mainWith $ do
update :: Weigh ()
update =
- let !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (const (E.fixed 0.5))
- :: M.Config ()
- !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.agrapa
- :: M.Config B.AGRAPA
- !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.ons
- :: M.Config B.ONS
+ let !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (E.Fixed 0.5)
+ !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.Agrapa
+ !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
!stF = M.initial cfgF
!stA = M.initial cfgA
!stO = M.initial cfgO
@@ -42,7 +36,7 @@ update =
decide :: Weigh ()
decide =
- let !cfg = M.config 0.5 0.0 1.0 1.0e-3 E.ons :: M.Config B.ONS
+ let !cfg = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
!st = M.initial cfg
in wgroup "Mean.decide" $ do
func "initial state" (M.decide cfg) st
@@ -50,12 +44,9 @@ decide =
stream :: Weigh ()
stream =
let !xs = force (take 1000 (cycle [0.3, 0.7]))
- !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (const (E.fixed 0.5))
- :: M.Config ()
- !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.agrapa
- :: M.Config B.AGRAPA
- !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.ons
- :: M.Config B.ONS
+ !cfgF = M.config 0.5 0.0 1.0 1.0e-3 (E.Fixed 0.5)
+ !cfgA = M.config 0.5 0.0 1.0 1.0e-3 E.Agrapa
+ !cfgO = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
runM cfg = foldl' (M.update cfg) (M.initial cfg)
in wgroup "Mean.update (1000-sample fold)" $ do
func "fixed" (runM cfgF) xs
@@ -65,12 +56,9 @@ stream =
twosample :: Weigh ()
twosample =
let !ps = force (take 1000 (cycle [(0.3, 0.7), (0.7, 0.3)]))
- !cfgF = TS.config 0.0 1.0 1.0e-3 (const (E.fixed 0.5))
- :: TS.Config ()
- !cfgA = TS.config 0.0 1.0 1.0e-3 E.agrapa
- :: TS.Config B.AGRAPA
- !cfgO = TS.config 0.0 1.0 1.0e-3 E.ons
- :: TS.Config B.ONS
+ !cfgF = TS.config 0.0 1.0 1.0e-3 (E.Fixed 0.5)
+ !cfgA = TS.config 0.0 1.0 1.0e-3 E.Agrapa
+ !cfgO = TS.config 0.0 1.0 1.0e-3 E.Ons
runT cfg = foldl' (TS.update cfg) (TS.initial cfg)
in wgroup "TwoSample.update (1000-sample fold)" $ do
func "fixed" (runT cfgF) ps
diff --git a/lib/Statistics/EProcess.hs b/lib/Statistics/EProcess.hs
@@ -28,12 +28,7 @@
module Statistics.EProcess (
-- * Bettors
- Bettor
- , AGRAPA
- , ONS
- , fixed
- , agrapa
- , ons
+ Bettor(..)
-- * Bounded-mean test
--
@@ -67,24 +62,24 @@ import qualified Statistics.EProcess.TwoSample as TS
-- | See 'Mean.config'.
meanConfig
- :: Double -- ^ null mean @m@
- -> Double -- ^ sample lower bound
- -> Double -- ^ sample upper bound
- -> Double -- ^ significance level @alpha@
- -> (Double -> Bettor s)
- -> Mean.Config s
+ :: Double -- ^ null mean @m@
+ -> Double -- ^ sample lower bound
+ -> Double -- ^ sample upper bound
+ -> Double -- ^ significance level @alpha@
+ -> Bettor
+ -> Mean.Config
meanConfig = Mean.config
-- | See 'Mean.initial'.
-initMeanState :: Mean.Config s -> Mean.State s
+initMeanState :: Mean.Config -> Mean.State
initMeanState = Mean.initial
-- | See 'Mean.update'.
-updateMean :: Mean.Config s -> Mean.State s -> Double -> Mean.State s
+updateMean :: Mean.Config -> Mean.State -> Double -> Mean.State
updateMean = Mean.update
-- | See 'Mean.decide'.
-decideMean :: Mean.Config s -> Mean.State s -> Mean.Verdict
+decideMean :: Mean.Config -> Mean.State -> Mean.Verdict
decideMean = Mean.decide
-- $twosample
@@ -97,27 +92,27 @@ twoSampleConfig
:: Double
-> Double
-> Double
- -> (Double -> Bettor s)
- -> TS.Config s
+ -> Bettor
+ -> TS.Config
twoSampleConfig = TS.config
-- | See 'TS.initial'.
-initTwoSampleState :: TS.Config s -> TS.State s
+initTwoSampleState :: TS.Config -> TS.State
initTwoSampleState = TS.initial
-- | See 'TS.update'.
updateTwoSample
- :: TS.Config s -> TS.State s -> (Double, Double) -> TS.State s
+ :: TS.Config -> TS.State -> (Double, Double) -> TS.State
updateTwoSample = TS.update
-- | See 'TS.decide'.
-decideTwoSample :: TS.Config s -> TS.State s -> TS.Verdict
+decideTwoSample :: TS.Config -> TS.State -> TS.Verdict
decideTwoSample = TS.decide
-- | Current log-wealth of a 'Mean.State'.
-logWealth :: Mean.State s -> Double
+logWealth :: Mean.State -> Double
logWealth = Mean.logWealth
-- | Sample count consumed so far.
-samples :: Mean.State s -> Int
+samples :: Mean.State -> Int
samples = Mean.samples
diff --git a/lib/Statistics/EProcess/Bettor.hs b/lib/Statistics/EProcess/Bettor.hs
@@ -1,6 +1,5 @@
{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
-{-# LANGUAGE RecordWildCards #-}
-- |
-- Module: Statistics.EProcess.Bettor
@@ -8,9 +7,9 @@
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
--- Bettor strategies for the e-process framework. A bettor maintains
--- internal state, consumes centred observations @z = x - m@, and
--- produces a predictable bet @lambda@ for the next observation.
+-- Bettor strategies for the e-process framework. A bettor describes
+-- how, given the history of centred observations @z = x - m@, the
+-- next predictable bet @lambda@ is chosen.
--
-- The bet placed at step @t@ depends only on data observed through
-- step @t-1@; this predictability is what makes the resulting wealth
@@ -18,106 +17,31 @@
-- and hence anytime-valid via Ville's inequality.
module Statistics.EProcess.Bettor (
- -- * Bettor type
+ -- * Bettor strategies
Bettor(..)
-
- -- * Strategies
- , fixed
- , agrapa
- , ons
-
- -- * Strategy state types (opaque)
- , AGRAPA
- , ONS
) where
-- | A predictable bettor.
--
--- Parameterised over its internal state type @s@.
---
--- * @bettorInit@: initial state.
--- * @bettorStep@: update state with a newly observed centred
--- value @z = x - m@.
--- * @bettorBet@: bet @lambda@ to use for the /next/ observation,
--- given the current state.
-data Bettor s = Bettor
- { bettorInit :: !s
- , bettorStep :: !(s -> Double -> s)
- , bettorBet :: !(s -> Double)
- }
-
--- | Fixed-lambda bettor.
---
--- Always bets the same value. Useful for smoke-testing the
--- framework and as a numerical baseline.
-fixed :: Double -> Bettor ()
-fixed !lam = Bettor
- { bettorInit = ()
- , bettorStep = \_ _ -> ()
- , bettorBet = \_ -> lam
- }
-
--- | aGRAPA bettor state (opaque).
-data AGRAPA = AGRAPA
- { aSum :: {-# UNPACK #-} !Double
- , aSum2 :: {-# UNPACK #-} !Double
- , aN :: {-# UNPACK #-} !Int
- , aMax :: {-# UNPACK #-} !Double
- }
-
--- | aGRAPA (approximate growth-rate adaptive predictable plug-in).
---
--- Tracks empirical mean and variance of centred observations @z@,
--- and bets the Kelly-optimal @lambda* = mu_z / (sigma_z^2 + mu_z^2)@
--- given the current point estimate, clipped to @[0, lambda_max]@.
---
--- The argument is @lambda_max@, the largest safe bet. For
--- observations @z = x - m@ where @x@ lies in @[lo, hi]@ and we are
--- testing @E[x] <= m@, a safe choice is @lambda_max = 1 \/ (m - lo)@
--- (so that the wealth factor @1 + lambda * z@ stays nonnegative).
-agrapa :: Double -> Bettor AGRAPA
-agrapa !lamMax = Bettor
- { bettorInit = AGRAPA 0 0 0 lamMax
- , bettorStep = \AGRAPA{..} !z ->
- AGRAPA (aSum + z) (aSum2 + z * z) (aN + 1) aMax
- , bettorBet = \AGRAPA{..} ->
- if aN == 0
- then 0
- else
- let !n = fromIntegral aN
- !mu = aSum / n
- !mu2 = mu * mu
- !var = max 0 (aSum2 / n - mu2)
- !den = var + mu2
- !raw = if den == 0 then 0 else mu / den
- in max 0 (min aMax raw)
- }
-
--- | ONS bettor state (opaque).
-data ONS = ONS
- { onsLambda :: {-# UNPACK #-} !Double
- , onsAcc :: {-# UNPACK #-} !Double
- , onsMax :: {-# UNPACK #-} !Double
- }
-
--- | ONS (online Newton step) bettor.
---
--- Maintains a running sum of squared gradients of the per-step
--- log-wealth loss and updates @lambda@ by a Newton step at each
--- observation. Achieves logarithmic regret against the best
--- constant bet in hindsight; in practice the strongest of the
--- three bettors here under most signal regimes.
---
--- The argument is @lambda_max@; see 'agrapa' for the sizing rule.
-ons :: Double -> Bettor ONS
-ons !lamMax = Bettor
- { bettorInit = ONS 0 1.0e-6 lamMax
- , bettorStep = \ONS{..} !z ->
- let !denom = 1 + onsLambda * z
- !g = if denom == 0 then 0 else negate z / denom
- !acc' = onsAcc + g * g
- !lam' = onsLambda - g / acc'
- !clp = max 0 (min onsMax lam')
- in ONS clp acc' onsMax
- , bettorBet = onsLambda
- }
+-- * 'Fixed' always bets the supplied @lambda@; useful for
+-- smoke-testing the framework and as a numerical baseline.
+--
+-- * 'Agrapa' is the aGRAPA (approximate growth-rate adaptive
+-- predictable plug-in) bettor; tracks empirical mean and variance
+-- of centred observations and bets the Kelly-optimal value given
+-- the current point estimate, clipped to @[0, lambda_max]@.
+--
+-- * 'Ons' is the online Newton step bettor; maintains a running
+-- sum of squared gradients of the per-step log-wealth loss and
+-- updates @lambda@ by a Newton step at each observation. Achieves
+-- logarithmic regret against the best constant bet in hindsight,
+-- and is in practice the strongest of the three under most signal
+-- regimes.
+--
+-- For 'Agrapa' and 'Ons', @lambda_max@ is derived from the sample
+-- bounds supplied to the surrounding test 'Statistics.EProcess.Mean.config'.
+data Bettor
+ = Fixed {-# UNPACK #-} !Double
+ | Agrapa
+ | Ons
+ deriving (Eq, Show)
diff --git a/lib/Statistics/EProcess/Mean.hs b/lib/Statistics/EProcess/Mean.hs
@@ -1,5 +1,6 @@
{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RecordWildCards #-}
-- |
@@ -37,16 +38,31 @@ module Statistics.EProcess.Mean (
, samples
) where
+import GHC.Exts (Double(D#))
import Statistics.EProcess.Bettor
-- | Test outcome at the current sample count.
data Verdict = Reject | Continue
deriving (Eq, Show)
+-- per-direction bettor state. one constructor per 'Bettor'
+-- alternative; the constructor used in a given 'State' matches the
+-- 'Bettor' chosen in the surrounding 'Config'.
+data BetState
+ = SFixed
+ | SAgrapa
+ {-# UNPACK #-} !Double -- sum of z
+ {-# UNPACK #-} !Double -- sum of z^2
+ {-# UNPACK #-} !Int -- count
+ | SOns
+ {-# UNPACK #-} !Double -- lambda
+ {-# UNPACK #-} !Double -- acc (sum of squared gradients)
+
-- | Test configuration. Constructed by 'config'.
-data Config s = Config
- { cfgBetPos :: !(Bettor s)
- , cfgBetNeg :: !(Bettor s)
+data Config = Config
+ { cfgBettor :: !Bettor
+ , cfgLamMaxPos :: {-# UNPACK #-} !Double
+ , cfgLamMaxNeg :: {-# UNPACK #-} !Double
, cfgNullMean :: {-# UNPACK #-} !Double
, cfgAlpha :: {-# UNPACK #-} !Double
, cfgLogThresh :: {-# UNPACK #-} !Double
@@ -54,78 +70,141 @@ data Config s = Config
-- | Test state. Two log-wealth processes (one per direction) and
-- per-direction bettor state.
-data State s = State
+data State = State
{ stN :: {-# UNPACK #-} !Int
, stLogWPos :: {-# UNPACK #-} !Double
, stLogWNeg :: {-# UNPACK #-} !Double
- , stBetPos :: !s
- , stBetNeg :: !s
+ , stBetPos :: !BetState
+ , stBetNeg :: !BetState
}
+-- floor for the wealth factor before taking a log; keeps the running
+-- log-wealth finite when a step pushes the factor to (or below) zero.
+-- NB. written via MagicHash because the fractional literal '1.0e-300'
+-- compiles as 'fromRational (1.0e-300 :: Rational)', and GHC does
+-- not constant-fold the conversion -- leaving a per-step
+-- '$wrationalToDouble' call in the worker.
+tiny :: Double
+tiny = D# 1.0e-300##
+{-# INLINE tiny #-}
+
-- | Build a test configuration.
--
--- The bettor argument is a function from @lambda_max@ to a bettor;
--- the same builder is used for both directions, with appropriate
--- bounds computed from @lo@, @hi@, and @m@.
---
-- >>> import qualified Statistics.EProcess.Bettor as B
--- >>> let cfg = config 0.5 0.0 1.0 1.0e-6 B.ons
+-- >>> let cfg = config 0.5 0.0 1.0 1.0e-6 B.Ons
config
- :: Double -- ^ null mean @m@
- -> Double -- ^ sample lower bound @lo@
- -> Double -- ^ sample upper bound @hi@
- -> Double -- ^ significance level @alpha@
- -> (Double -> Bettor s) -- ^ bettor builder, taking @lambda_max@
- -> Config s
-config !m !lo !hi !alpha mk = Config
- { cfgBetPos = mk (0.5 / (m - lo))
- , cfgBetNeg = mk (0.5 / (hi - m))
+ :: Double -- ^ null mean @m@
+ -> Double -- ^ sample lower bound @lo@
+ -> Double -- ^ sample upper bound @hi@
+ -> Double -- ^ significance level @alpha@
+ -> Bettor -- ^ bettor strategy
+ -> Config
+config !m !lo !hi !alpha !b = Config
+ { cfgBettor = b
+ , cfgLamMaxPos = 0.5 / (m - lo)
+ , cfgLamMaxNeg = 0.5 / (hi - m)
, cfgNullMean = m
, cfgAlpha = alpha
, cfgLogThresh = log (2 / alpha)
}
--- NB. argument to @mk@ is half the geometric @lambda_max@; the 1/2
+-- NB. the lambda_max values are half the geometric ceiling; the 1/2
-- margin keeps the wealth factor bounded away from zero at the
-- boundary, which is the WSR safety recommendation.
+{-# INLINE config #-}
+
+-- per-bettor initial state.
+initBet :: Bettor -> BetState
+initBet b = case b of
+ Fixed _ -> SFixed
+ Agrapa -> SAgrapa 0 0 0
+ Ons -> SOns 0 1.0e-6
+{-# INLINE initBet #-}
-- | Initial state for streaming.
-initial :: Config s -> State s
-initial Config{..} = State
- { stN = 0
- , stLogWPos = 0
- , stLogWNeg = 0
- , stBetPos = bettorInit cfgBetPos
- , stBetNeg = bettorInit cfgBetNeg
- }
+initial :: Config -> State
+initial Config{..} =
+ let !s0 = initBet cfgBettor
+ in State
+ { stN = 0
+ , stLogWPos = 0
+ , stLogWNeg = 0
+ , stBetPos = s0
+ , stBetNeg = s0
+ }
+{-# INLINE initial #-}
+
+-- compute the next bet @lambda@ from the bettor and its current
+-- state. @lamMax@ is the direction-specific safety bound.
+betLambda :: Bettor -> Double -> BetState -> Double
+betLambda b !lamMax !s = case b of
+ Fixed lam -> lam
+ Agrapa -> case s of
+ SAgrapa !sm !sm2 !n
+ | n == 0 -> 0
+ | otherwise ->
+ let !nd = fromIntegral n
+ !mu = sm / nd
+ !mu2 = mu * mu
+ !var = max 0 (sm2 / nd - mu2)
+ !den = var + mu2
+ !raw = if den == 0 then 0 else mu / den
+ in max 0 (min lamMax raw)
+ _ -> 0
+ Ons -> case s of
+ SOns !lam _ -> lam
+ _ -> 0
+{-# INLINE betLambda #-}
+
+-- update bettor state with newly observed centred value @z@.
+stepBet :: Bettor -> Double -> BetState -> Double -> BetState
+stepBet b !lamMax !s !z = case b of
+ Fixed _ -> SFixed
+ Agrapa -> case s of
+ SAgrapa !sm !sm2 !n -> SAgrapa (sm + z) (sm2 + z * z) (n + 1)
+ _ -> SAgrapa z (z * z) 1
+ Ons -> case s of
+ SOns !lam !acc ->
+ let !denom = 1 + lam * z
+ !g = if denom == 0 then 0 else negate z / denom
+ !acc' = acc + g * g
+ !lam' = lam - g / acc'
+ !clp = max 0 (min lamMax lam')
+ in SOns clp acc'
+ _ -> SOns 0 1.0e-6
+{-# INLINE stepBet #-}
-- | Fold one observation into the state.
-update :: Config s -> State s -> Double -> State s
+update :: Config -> State -> Double -> State
update Config{..} State{..} !x =
let !z = x - cfgNullMean
- !lamP = bettorBet cfgBetPos stBetPos
- !lamN = bettorBet cfgBetNeg stBetNeg
+ !lamP = betLambda cfgBettor cfgLamMaxPos stBetPos
+ !lamN = betLambda cfgBettor cfgLamMaxNeg stBetNeg
!facP = 1 + lamP * z
!facN = 1 - lamN * z
- !logWP' = stLogWPos + log (max 1.0e-300 facP)
- !logWN' = stLogWNeg + log (max 1.0e-300 facN)
- !sP' = bettorStep cfgBetPos stBetPos z
- !sN' = bettorStep cfgBetNeg stBetNeg (negate z)
+ !logWP' = stLogWPos + log (max tiny facP)
+ !logWN' = stLogWNeg + log (max tiny facN)
+ !sP' = stepBet cfgBettor cfgLamMaxPos stBetPos z
+ !sN' = stepBet cfgBettor cfgLamMaxNeg stBetNeg (negate z)
in State (stN + 1) logWP' logWN' sP' sN'
+{-# INLINE update #-}
-- | Decide based on current wealth.
--
-- 'Reject' iff either directional log-wealth has crossed the
-- Bonferroni-adjusted threshold @log(2 \/ alpha)@.
-decide :: Config s -> State s -> Verdict
+decide :: Config -> State -> Verdict
decide Config{..} State{..}
| stLogWPos >= cfgLogThresh = Reject
| stLogWNeg >= cfgLogThresh = Reject
| otherwise = Continue
+{-# INLINE decide #-}
-- | Current log-wealth (the larger of the two directional processes).
-logWealth :: State s -> Double
+logWealth :: State -> Double
logWealth State{..} = max stLogWPos stLogWNeg
+{-# INLINE logWealth #-}
-- | Sample count consumed so far.
-samples :: State s -> Int
+samples :: State -> Int
samples = stN
+{-# INLINE samples #-}
diff --git a/lib/Statistics/EProcess/TwoSample.hs b/lib/Statistics/EProcess/TwoSample.hs
@@ -38,10 +38,10 @@ import Statistics.EProcess.Mean (Verdict(..))
import Statistics.EProcess.Bettor (Bettor)
-- | Test configuration.
-newtype Config s = Config (M.Config s)
+newtype Config = Config M.Config
-- | Test state.
-newtype State s = State (M.State s)
+newtype State = State M.State
-- | Build a paired two-sample test configuration.
--
@@ -49,34 +49,40 @@ newtype State s = State (M.State s)
-- samples; differences then lie in @[lo - hi, hi - lo]@.
--
-- >>> import qualified Statistics.EProcess.Bettor as B
--- >>> let cfg = config 0.0 1.0 1.0e-6 B.ons
+-- >>> let cfg = config 0.0 1.0 1.0e-6 B.Ons
config
- :: Double -- ^ sample lower bound @lo@
- -> Double -- ^ sample upper bound @hi@
- -> Double -- ^ significance level @alpha@
- -> (Double -> Bettor s) -- ^ bettor builder
- -> Config s
-config !lo !hi !alpha mk =
- let !b = hi - lo
- in Config (M.config 0 (negate b) b alpha mk)
+ :: Double -- ^ sample lower bound @lo@
+ -> Double -- ^ sample upper bound @hi@
+ -> Double -- ^ significance level @alpha@
+ -> Bettor -- ^ bettor strategy
+ -> Config
+config !lo !hi !alpha b =
+ let !d = hi - lo
+ in Config (M.config 0 (negate d) d alpha b)
+{-# INLINE config #-}
-- | Initial state for streaming.
-initial :: Config s -> State s
+initial :: Config -> State
initial (Config c) = State (M.initial c)
+{-# INLINE initial #-}
-- | Fold one paired observation @(a, b)@ into the state.
-update :: Config s -> State s -> (Double, Double) -> State s
+update :: Config -> State -> (Double, Double) -> State
update (Config c) (State s) (!a, !b) =
State (M.update c s (a - b))
+{-# INLINE update #-}
-- | Decide based on current wealth.
-decide :: Config s -> State s -> Verdict
+decide :: Config -> State -> Verdict
decide (Config c) (State s) = M.decide c s
+{-# INLINE decide #-}
-- | Current log-wealth.
-logWealth :: State s -> Double
+logWealth :: State -> Double
logWealth (State s) = M.logWealth s
+{-# INLINE logWealth #-}
-- | Sample count consumed so far.
-samples :: State s -> Int
+samples :: State -> Int
samples (State s) = M.samples s
+{-# INLINE samples #-}
diff --git a/test/Main.hs b/test/Main.hs
@@ -48,7 +48,7 @@ bernoulli !p g =
-- with the early-stopping rule built in. returns (verdict, samples
-- consumed).
runMeanBernoulli
- :: M.Config s
+ :: M.Config
-> Double -- ^ p
-> Int -- ^ budget
-> Gen
@@ -66,7 +66,7 @@ runMeanBernoulli cfg p budget g0 = go 0 g0 (M.initial cfg)
-- fraction of trials that rejected.
rejectionRate
- :: M.Config s
+ :: M.Config
-> Double -- ^ true bernoulli p
-> Int -- ^ budget per trial
-> Int -- ^ number of trials
@@ -88,12 +88,12 @@ genSeq g = let (_, g') = stepGen g in g : genSeq g'
sanityTests :: TestTree
sanityTests = testGroup "sanity" [
testCase "degenerate input never rejects" $ do
- let cfg = M.config 0.5 0.0 1.0 1.0e-6 E.ons :: M.Config E.ONS
+ let cfg = M.config 0.5 0.0 1.0 1.0e-6 E.Ons
xs = replicate 5000 0.5
st = foldl' (M.update cfg) (M.initial cfg) xs
M.decide cfg st @?= M.Continue
, testCase "two-sided thresholds applied symmetrically" $ do
- let cfg = M.config 0.5 0.0 1.0 1.0e-6 E.ons :: M.Config E.ONS
+ let cfg = M.config 0.5 0.0 1.0 1.0e-6 E.Ons
M.decide cfg (M.initial cfg) @?= M.Continue
]
@@ -104,14 +104,14 @@ sanityTests = testGroup "sanity" [
calibrationTests :: TestTree
calibrationTests = testGroup "null calibration" [
testCase "ONS, Bernoulli(0.5), m=0.5, alpha=0.05" $ do
- let cfg = M.config 0.5 0.0 1.0 0.05 E.ons :: M.Config E.ONS
+ let cfg = M.config 0.5 0.0 1.0 0.05 E.Ons
rate = rejectionRate cfg 0.5 2000 200 12345
-- expected rate ≤ 0.05; allow up to 0.10 slack for sampling
-- variability over 200 trials.
assertBool ("FPR " ++ show rate ++ " exceeded slack") $
rate <= 0.10
, testCase "aGRAPA, Bernoulli(0.5), m=0.5, alpha=0.05" $ do
- let cfg = M.config 0.5 0.0 1.0 0.05 E.agrapa :: M.Config E.AGRAPA
+ let cfg = M.config 0.5 0.0 1.0 0.05 E.Agrapa
rate = rejectionRate cfg 0.5 2000 200 67890
assertBool ("FPR " ++ show rate ++ " exceeded slack") $
rate <= 0.10
@@ -123,13 +123,12 @@ calibrationTests = testGroup "null calibration" [
powerTests :: TestTree
powerTests = testGroup "power" [
testCase "ONS detects Bernoulli(0.7) vs m=0.5" $ do
- let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.ons :: M.Config E.ONS
+ let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
rate = rejectionRate cfg 0.7 5000 100 11111
assertBool ("power " ++ show rate ++ " too low") $
rate >= 0.95
, testCase "aGRAPA detects Bernoulli(0.7) vs m=0.5" $ do
- let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.agrapa
- :: M.Config E.AGRAPA
+ let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.Agrapa
rate = rejectionRate cfg 0.7 5000 100 22222
assertBool ("power " ++ show rate ++ " too low") $
rate >= 0.95
@@ -138,7 +137,7 @@ powerTests = testGroup "power" [
-- two-sample paired test.
runTSPaired
- :: TS.Config s
+ :: TS.Config
-> Double
-> Double -- ^ p for A and B
-> Int
@@ -159,11 +158,11 @@ runTSPaired cfg pA pB budget g0 = go 0 g0 (TS.initial cfg)
twoSampleTests :: TestTree
twoSampleTests = testGroup "two-sample" [
testCase "identical distributions don't reject" $ do
- let cfg = TS.config 0.0 1.0 1.0e-3 E.ons :: TS.Config E.ONS
+ let cfg = TS.config 0.0 1.0 1.0e-3 E.Ons
rate = avgRate cfg 0.5 0.5 2000 100 33333
assertBool ("FPR " ++ show rate) $ rate <= 0.05
, testCase "different distributions reject" $ do
- let cfg = TS.config 0.0 1.0 1.0e-3 E.ons :: TS.Config E.ONS
+ let cfg = TS.config 0.0 1.0 1.0e-3 E.Ons
rate = avgRate cfg 0.3 0.7 5000 100 44444
assertBool ("power " ++ show rate) $ rate >= 0.95
]
@@ -182,19 +181,17 @@ twoSampleTests = testGroup "two-sample" [
bettorSmokeTests :: TestTree
bettorSmokeTests = testGroup "bettor smoke" [
testCase "fixed bettor runs without error" $ do
- let cfg = M.config 0.5 0.0 1.0 1.0e-3
- (const (E.fixed 0.5)) :: M.Config ()
+ let cfg = M.config 0.5 0.0 1.0 1.0e-3 (E.Fixed 0.5)
xs = take 100 (cycle [0.0, 1.0])
st = foldl' (M.update cfg) (M.initial cfg) xs
assertBool "samples advanced" (M.samples st == 100)
, testCase "ONS bettor runs without error" $ do
- let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.ons :: M.Config E.ONS
+ let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.Ons
xs = take 100 (cycle [0.0, 1.0])
st = foldl' (M.update cfg) (M.initial cfg) xs
assertBool "samples advanced" (M.samples st == 100)
, testCase "aGRAPA bettor runs without error" $ do
- let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.agrapa
- :: M.Config E.AGRAPA
+ let cfg = M.config 0.5 0.0 1.0 1.0e-3 E.Agrapa
xs = take 100 (cycle [0.0, 1.0])
st = foldl' (M.update cfg) (M.initial cfg) xs
assertBool "samples advanced" (M.samples st == 100)
