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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
|
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TupleSections #-}
module PPL.Sampling
( mh,
ssmh,
)
where
import Control.DeepSeq
import Control.Exception (evaluate)
import Control.Monad.IO.Class
import Control.Monad.Trans.State
import Data.Bifunctor
import qualified Data.Map.Strict as M
import Data.Monoid
import GHC.Exts.Heap
import Numeric.Log
import PPL.Distr
import PPL.Internal
import qualified Streaming as S
import Streaming.Prelude (Of, Stream, yield)
import System.IO.Unsafe
import System.Random (StdGen, random, randoms)
import qualified System.Random as R
import Data.IORef
import Control.Monad
import Debug.Trace
mh :: (MonadIO m) => StdGen -> Double -> Meas a -> Stream (Of (a, Log Double)) m ()
mh g p m = do
let (g0, g1) = R.split g
omega <- liftIO $ newIORef (mempty, 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 (M.Map [Int] Double, StdGen) -> m (IORef (M.Map [Int] Double, StdGen))
mutate g omega = do
(m, g0) <- liftIO $ readIORef omega
let (r:q:_) = R.randoms g
ks = M.keys m
k = ks !! floor (r * join traceShow (fromIntegral (length ks)))
m' = M.insert k q m
liftIO $ newIORef $ (m',g0)
where
go x = do
g <- get
let (r, g1) = R.random g
(y, g2) = R.random g1
if r < p
then do
put g2
pure y
else do
put g1
pure x
-- Single site MH
-- Truncated trees
data TTree = TTree Bool [Maybe TTree] deriving (Show)
type Site = [Int]
trunc :: Tree -> IO TTree
trunc = truncTree . asBox
where
truncTree t =
getBoxedClosureData' t >>= \case
ConstrClosure _ [l, r] [] _ _ "Tree" ->
getBoxedClosureData' l >>= \case
ConstrClosure {dataArgs = [d], name = "D#"} ->
TTree True <$> truncTrees r
x -> error $ "truncTree:ConstrClosure:" ++ show x
ConstrClosure _ [r] [d] _ _ "Tree" ->
TTree False <$> truncTrees r
x -> error $ "truncTree:" ++ show x
getBoxedClosureData' x =
getBoxedClosureData x >>= \c -> case c of
BlackholeClosure _ t -> getBoxedClosureData' t
_ -> pure c
truncTrees b =
getBoxedClosureData' b >>= \case
ConstrClosure _ [l, r] [] _ _ ":" ->
getBoxedClosureData' l >>= \case
ConstrClosure {name = "Tree"} ->
((:) . Just) <$> truncTree l <*> truncTrees r
_ -> (Nothing :) <$> truncTrees r
_ -> pure []
trunc' t x w = unsafePerformIO $ do
evaluate (rnf x)
evaluate (rnf w)
trunc t
sites :: Site -> TTree -> [Site]
sites acc (TTree eval ts) = (if eval then acc else mempty) : concat [sites (x : acc) t | (x, Just t) <- zip [0 ..] ts]
mutate = M.foldrWithKey go
where
go [] d (Tree _ ts) = Tree d ts
go (n : ns) d (Tree v ts) = Tree v $ take n ts ++ go ns d (ts !! n) : drop (n + 1) ts
ssmh :: (Show a, NFData a, Monad m) => StdGen -> Meas a -> Stream (Of (a, Log Double)) m ()
ssmh g m = step t (mempty :: M.Map Site Double) (trunc' t0 x w) x w
where
(t0, t) = split $ randomTree g
(x, w) = head $ samples m t0
step !t !sub !tt !x !w = do
let ss = sites [] tt
(t1 : t2 : t3 : t4 : _) = splitTrees t
i = floor $ draw t2 * (fromIntegral $ length ss) -- site to mutate
sub' = M.insert (reverse $ ss !! i) (draw t3) sub
t' = mutate t0 sub'
(x', w') = head $ samples m t'
tt' = trunc' t' x' w'
ratio = w' / w
(Exp . log -> r) = draw t4
(sub'', tt'', x'', w'') =
if r < ratio
then (sub', tt', x', w')
else (sub, tt, x, w)
yield (x'', w'')
step t1 sub'' tt'' x'' w''
|
