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
|
{-# 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
import PPL.Internal
import Streaming.Prelude (Of, Stream, yield)
import System.Random (StdGen)
import System.Random qualified as R
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) = R.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, R.split -> (g1, g2)) = R.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]) = (R.randoms *** R.randoms) (R.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)
|
