{-# LANGUAGE RankNTypes #-}
module Control.Monad.Bayes.Traced.Basic
( TracedT,
hoist,
marginal,
mhStep,
mh,
)
where
import Control.Applicative (liftA2)
import Control.Monad.Bayes.Class
( MonadDistribution (random),
MonadFactor (..),
MonadMeasure,
)
import Control.Monad.Bayes.Density.Free (DensityT)
import Control.Monad.Bayes.Traced.Common
( Trace (..),
bind,
mhTrans',
scored,
singleton,
)
import Control.Monad.Bayes.Weighted (WeightedT)
import Data.Functor.Identity (Identity)
import Data.List.NonEmpty as NE (NonEmpty ((:|)), toList)
data TracedT m a = TracedT
{
forall (m :: * -> *) a.
TracedT m a -> WeightedT (DensityT Identity) a
model :: WeightedT (DensityT Identity) a,
forall (m :: * -> *) a. TracedT m a -> m (Trace a)
traceDist :: m (Trace a)
}
instance (Monad m) => Functor (TracedT m) where
fmap :: forall a b. (a -> b) -> TracedT m a -> TracedT m b
fmap a -> b
f (TracedT WeightedT (DensityT Identity) a
m m (Trace a)
d) = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f WeightedT (DensityT Identity) a
m) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) m (Trace a)
d)
instance (Monad m) => Applicative (TracedT m) where
pure :: forall a. a -> TracedT m a
pure a
x = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x) (forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x))
(TracedT WeightedT (DensityT Identity) (a -> b)
mf m (Trace (a -> b))
df) <*> :: forall a b. TracedT m (a -> b) -> TracedT m a -> TracedT m b
<*> (TracedT WeightedT (DensityT Identity) a
mx m (Trace a)
dx) = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT (WeightedT (DensityT Identity) (a -> b)
mf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> WeightedT (DensityT Identity) a
mx) (forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) m (Trace (a -> b))
df m (Trace a)
dx)
instance (Monad m) => Monad (TracedT m) where
(TracedT WeightedT (DensityT Identity) a
mx m (Trace a)
dx) >>= :: forall a b. TracedT m a -> (a -> TracedT m b) -> TracedT m b
>>= a -> TracedT m b
f = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT WeightedT (DensityT Identity) b
my m (Trace b)
dy
where
my :: WeightedT (DensityT Identity) b
my = WeightedT (DensityT Identity) a
mx forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a.
TracedT m a -> WeightedT (DensityT Identity) a
model forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> TracedT m b
f
dy :: m (Trace b)
dy = m (Trace a)
dx forall (m :: * -> *) a b.
Monad m =>
m (Trace a) -> (a -> m (Trace b)) -> m (Trace b)
`bind` (forall (m :: * -> *) a. TracedT m a -> m (Trace a)
traceDist forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> TracedT m b
f)
instance (MonadDistribution m) => MonadDistribution (TracedT m) where
random :: TracedT m Double
random = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT forall (m :: * -> *). MonadDistribution m => m Double
random (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Double -> Trace Double
singleton forall (m :: * -> *). MonadDistribution m => m Double
random)
instance (MonadFactor m) => MonadFactor (TracedT m) where
score :: Log Double -> TracedT m ()
score Log Double
w = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT (forall (m :: * -> *). MonadFactor m => Log Double -> m ()
score Log Double
w) (forall (m :: * -> *). MonadFactor m => Log Double -> m ()
score Log Double
w forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Log Double -> Trace ()
scored Log Double
w))
instance (MonadMeasure m) => MonadMeasure (TracedT m)
hoist :: (forall x. m x -> m x) -> TracedT m a -> TracedT m a
hoist :: forall (m :: * -> *) a.
(forall x. m x -> m x) -> TracedT m a -> TracedT m a
hoist forall x. m x -> m x
f (TracedT WeightedT (DensityT Identity) a
m m (Trace a)
d) = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT WeightedT (DensityT Identity) a
m (forall x. m x -> m x
f m (Trace a)
d)
marginal :: (Monad m) => TracedT m a -> m a
marginal :: forall (m :: * -> *) a. Monad m => TracedT m a -> m a
marginal (TracedT WeightedT (DensityT Identity) a
_ m (Trace a)
d) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Trace a -> a
output m (Trace a)
d
mhStep :: (MonadDistribution m) => TracedT m a -> TracedT m a
mhStep :: forall (m :: * -> *) a.
MonadDistribution m =>
TracedT m a -> TracedT m a
mhStep (TracedT WeightedT (DensityT Identity) a
m m (Trace a)
d) = forall (m :: * -> *) a.
WeightedT (DensityT Identity) a -> m (Trace a) -> TracedT m a
TracedT WeightedT (DensityT Identity) a
m m (Trace a)
d'
where
d' :: m (Trace a)
d' = m (Trace a)
d forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a.
MonadDistribution m =>
WeightedT (DensityT Identity) a -> Trace a -> m (Trace a)
mhTrans' WeightedT (DensityT Identity) a
m
mh :: (MonadDistribution m) => Int -> TracedT m a -> m [a]
mh :: forall (m :: * -> *) a.
MonadDistribution m =>
Int -> TracedT m a -> m [a]
mh Int
n (TracedT WeightedT (DensityT Identity) a
m m (Trace a)
d) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> [a] -> [b]
map forall a. Trace a -> a
output forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NonEmpty a -> [a]
NE.toList) (forall {t}. (Ord t, Num t) => t -> m (NonEmpty (Trace a))
f Int
n)
where
f :: t -> m (NonEmpty (Trace a))
f t
k
| t
k forall a. Ord a => a -> a -> Bool
<= t
0 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. a -> [a] -> NonEmpty a
:| []) m (Trace a)
d
| Bool
otherwise = do
(Trace a
x :| [Trace a]
xs) <- t -> m (NonEmpty (Trace a))
f (t
k forall a. Num a => a -> a -> a
- t
1)
Trace a
y <- forall (m :: * -> *) a.
MonadDistribution m =>
WeightedT (DensityT Identity) a -> Trace a -> m (Trace a)
mhTrans' WeightedT (DensityT Identity) a
m Trace a
x
forall (m :: * -> *) a. Monad m => a -> m a
return (Trace a
y forall a. a -> [a] -> NonEmpty a
:| Trace a
x forall a. a -> [a] -> [a]
: [Trace a]
xs)