1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
|
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
module PPL.Sampling
( mh,
)
where
import Control.Arrow
import Control.Monad
import Control.Monad.IO.Class
import Data.Bits (countTrailingZeros, shiftL, shiftR)
import Data.IORef
import Data.List (foldl')
import Data.Vector qualified as V
import Data.Vector.Hashtables qualified as H
import Data.Word (Word64)
import Numeric.Log hiding (sum)
import PPL.Internal hiding (newTree, split)
import Streaming.Prelude (Of, Stream, yield)
import System.IO.Unsafe
import System.Random (StdGen, split)
import System.Random qualified as R
import Unsafe.Coerce
type MemoTree = IORef (HashMap Integer Word64, StdGen)
-- Let's draw doubles more uniformly than the standard generator
-- based on https://www.corsix.org/content/higher-quality-random-floats
random :: StdGen -> (Double, StdGen)
random g0 =
let (r1 :: Word64, g1) = R.random g0
(r2 :: Word64, g2) = R.random g1
e = let r = countTrailingZeros r1 - 11 in if r >= 0 then countTrailingZeros r2 else r
dbl = (1 + (r1 `shiftR` 11)) `shiftR` 1 - ((unsafeCoerce e - 1011) `shiftL` 52)
in (unsafeCoerce dbl, g2)
{-# INLINE newMemoTree #-}
newMemoTree :: MemoTree -> Tree
newMemoTree s = go 1
where
go :: Integer -> Tree
go id =
Tree
( unsafePerformIO $ do
(m, g) <- readIORef s
H.lookup m id >>= \case
Nothing -> do
let (x, g') = R.random g
H.insert m id x
writeIORef s (m, g')
pure x
Just x -> pure x
)
(go (2 * id), go (2 * id + 1))
data MinHeap a k = Node a k (MinHeap a k) (MinHeap a k) | Empty
deriving (Show)
insert m k v = merge m (Node k v Empty Empty)
merge Empty n = n
merge n Empty = n
merge n0@(Node k0 v0 l0 r0) n1@(Node k1 v1 l1 r1) =
if k0 < k1
then Node k0 v0 (merge r0 n1) l0
else Node k1 v1 r1 (merge l1 n0)
toList Empty = []
toList (Node k v l r) = v : toList (merge l r)
mh :: (MonadIO m) => StdGen -> Double -> Meas a -> Stream (Of (a, Log Double)) m ()
mh g p m = do
let (g0, g1) = split g
hm <- liftIO $ H.initialize 0
omega <- liftIO $ newIORef (hm, g0)
let (x, w) = head $ samples m $ newMemoTree omega
step g1 omega x w
where
step !g0 !omega !x !w = do
let (Exp . log -> r, split -> (g1, g2)) = random g0
omega' <- mutate g1 omega
let (!x', !w') = head $ samples m $ newMemoTree omega'
ratio = w' / w
(omega'', x'', w'') =
if r < ratio
then (omega', x', w')
else (omega, x, w)
yield (x'', w'')
step g2 omega'' x'' w''
mutate :: (MonadIO m) => StdGen -> MemoTree -> m MemoTree
mutate g omega = liftIO $ do
(m, g0) <- readIORef omega
m' <- H.clone m
ks <- H.keys m
let (rs :: [Word64], qs :: [Word64]) = (R.randoms *** R.randoms) (split g)
ks' = toList $ foldl' (\m -> uncurry (insert m)) Empty $ zip rs $ V.toList ks
n = fromIntegral (V.length ks)
void $ zipWithM_ (\k q -> H.insert m' k q) (take (1 + floor (p * fromIntegral n)) ks') qs
newIORef (m', g0)
|
