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
|
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
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)
mh :: (Monad m) => StdGen -> Double -> Meas a -> Stream (Of (a, Log Double)) m ()
mh g p m = step t0 t x w
where
(t0, t) = split $ randomTree g
(x, w) = head $ samples m t
step !t0 !t !x !w = do
let (t1 : t2 : t3 : t4 : _) = splitTrees t0
t' = mutateTree p t1 t2 t
(x', w') = head $ samples m t'
ratio = w' / w
(Exp . log -> r) = draw t3
(t'', x'', w'') =
if r < ratio
then (t', x', w')
else (t, x, w)
yield (x'', w'')
step t4 t'' x'' w''
-- 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''
|
