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
|
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
module PPL.Internal
( uniform,
Prob (..),
Meas,
score,
scoreLog,
sample,
memoize,
samples,
newTree,
HashMap,
Tree (..),
)
where
import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Writer
import Data.Bifunctor
import Data.IORef
import Data.Map qualified as Q
import Data.Monoid
import Data.Vector.Hashtables qualified as H
import Data.Vector.Mutable qualified as M
import Data.Vector.Unboxed.Mutable qualified as UM
import Language.Haskell.TH.Syntax qualified as TH
import Numeric.Log
import System.IO.Unsafe
import System.Random hiding (split, uniform)
import System.Random qualified as R
type HashMap k v = H.Dictionary (H.PrimState IO) M.MVector k UM.MVector v
-- Reimplementation of the LazyPPL monads to avoid some dependencies
data Tree = Tree
{ draw :: Double,
split :: (Tree, Tree)
}
{-# 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') = 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))
newtype Prob a = Prob {runProb :: Tree -> a}
instance Monad Prob where
Prob f >>= g = Prob $ \t ->
let (t1, t2) = split t
(Prob g') = g (f t1)
in g' t2
instance Functor Prob where fmap = liftM
instance Applicative Prob where pure = Prob . const; (<*>) = ap
uniform :: Prob Double
uniform = Prob $ \(Tree r _) -> r
newtype Meas a = Meas (WriterT (Product (Log Double)) Prob a)
deriving (Functor, Applicative, Monad)
{-# INLINE score #-}
score :: Double -> Meas ()
score = scoreLog . Exp . log . max eps
where
eps = $(TH.lift (2 * until ((== 1) . (1 +)) (/ 2) (1 :: Double))) -- machine epsilon, force compile time eval
{-# INLINE scoreLog #-}
scoreLog :: Log Double -> Meas ()
scoreLog = Meas . tell . Product
{-# INLINE sample #-}
sample :: Prob a -> Meas a
sample = Meas . lift
{-# INLINE samples #-}
samples :: forall a. Meas a -> Tree -> [(a, Log Double)]
samples (Meas m) = map (second getProduct) . runProb f
where
f = runWriterT m >>= \x -> (x :) <$> f
{-# NOINLINE memoize #-}
memoize :: Ord a => (a -> Prob b) -> Prob (a -> b)
memoize f = unsafePerformIO $ do
ref <- newIORef mempty
pure $ Prob $ \t -> \x -> unsafePerformIO $ do
m <- readIORef ref
case Q.lookup x m of
Just z -> pure z
_ -> do
let Prob g = f x
z = g (getNode (Q.size m) t)
m' = Q.insert x z m
writeIORef ref m'
pure z
where
getNode 0 t = t
getNode i (split -> (t0, t1)) = getNode (i `div` 2) (if i `mod` 2 == 1 then t0 else t1)
|
