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
|
{-# 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.IORef
import Data.List (foldl')
import Data.Vector qualified as V
import Data.Vector.Hashtables qualified as H
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)
{-# INLINE newTree #-}
newTree :: IORef (HashMap Integer Double, StdGen) -> Tree
newTree 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') = 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 $ newTree 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 $ newTree 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 -> IORef (HashMap Integer Double, StdGen) -> m (IORef (HashMap Integer Double, StdGen))
mutate g omega = liftIO $ do
(m, g0) <- readIORef omega
m' <- H.clone m
ks <- H.keys m
let (rs :: [Double], qs :: [Double]) = (randoms *** 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)
|
