{-# LANGUAGE RankNTypes #-}

-- |
-- Module      : Control.Monad.Bayes.Traced.Dynamic
-- Description : Distributions on execution traces that can be dynamically frozen
-- Copyright   : (c) Adam Scibior, 2015-2020
-- License     : MIT
-- Maintainer  : leonhard.markert@tweag.io
-- Stability   : experimental
-- Portability : GHC
module Control.Monad.Bayes.Traced.Dynamic
  ( TracedT,
    hoist,
    marginal,
    freeze,
    mhStep,
    mh,
  )
where

import Control.Monad (join)
import Control.Monad.Bayes.Class
  ( MonadDistribution (random),
    MonadFactor (..),
    MonadMeasure,
  )
import Control.Monad.Bayes.Density.Free (DensityT)
import Control.Monad.Bayes.Traced.Common
  ( Trace (..),
    bind,
    mhTransFree,
    scored,
    singleton,
  )
import Control.Monad.Bayes.Weighted (WeightedT)
import Control.Monad.Trans (MonadTrans (..))
import Data.List.NonEmpty as NE (NonEmpty ((:|)), toList)

-- | A tracing monad where only a subset of random choices are traced and this
-- subset can be adjusted dynamically.
newtype TracedT m a = TracedT {forall (m :: * -> *) a.
TracedT m a -> m (WeightedT (DensityT m) a, Trace a)
runTraced :: m (WeightedT (DensityT m) a, Trace a)}

pushM :: (Monad m) => m (WeightedT (DensityT m) a) -> WeightedT (DensityT m) a
pushM :: forall (m :: * -> *) a.
Monad m =>
m (WeightedT (DensityT m) a) -> WeightedT (DensityT m) a
pushM = forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

instance (Monad m) => Functor (TracedT m) where
  fmap :: forall a b. (a -> b) -> TracedT m a -> TracedT m b
fmap a -> b
f (TracedT m (WeightedT (DensityT m) a, Trace a)
c) = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ do
    (WeightedT (DensityT m) a
m, Trace a
t) <- m (WeightedT (DensityT m) a, Trace a)
c
    let m' :: WeightedT (DensityT m) b
m' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f WeightedT (DensityT m) a
m
    let t' :: Trace b
t' = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Trace a
t
    forall (m :: * -> *) a. Monad m => a -> m a
return (WeightedT (DensityT m) b
m', Trace b
t')

instance (Monad m) => Applicative (TracedT m) where
  pure :: forall a. a -> TracedT m a
pure a
x = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x, forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)
  (TracedT m (WeightedT (DensityT m) (a -> b), Trace (a -> b))
cf) <*> :: forall a b. TracedT m (a -> b) -> TracedT m a -> TracedT m b
<*> (TracedT m (WeightedT (DensityT m) a, Trace a)
cx) = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ do
    (WeightedT (DensityT m) (a -> b)
mf, Trace (a -> b)
tf) <- m (WeightedT (DensityT m) (a -> b), Trace (a -> b))
cf
    (WeightedT (DensityT m) a
mx, Trace a
tx) <- m (WeightedT (DensityT m) a, Trace a)
cx
    forall (m :: * -> *) a. Monad m => a -> m a
return (WeightedT (DensityT m) (a -> b)
mf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> WeightedT (DensityT m) a
mx, Trace (a -> b)
tf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Trace a
tx)

instance (Monad m) => Monad (TracedT m) where
  (TracedT m (WeightedT (DensityT m) a, Trace a)
cx) >>= :: forall a b. TracedT m a -> (a -> TracedT m b) -> TracedT m b
>>= a -> TracedT m b
f = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ do
    (WeightedT (DensityT m) a
mx, Trace a
tx) <- m (WeightedT (DensityT m) a, Trace a)
cx
    let m :: WeightedT (DensityT m) b
m = WeightedT (DensityT m) a
mx forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a.
Monad m =>
m (WeightedT (DensityT m) a) -> WeightedT (DensityT m) a
pushM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
TracedT m a -> m (WeightedT (DensityT m) a, Trace a)
runTraced forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> TracedT m b
f
    Trace b
t <- forall (m :: * -> *) a. Monad m => a -> m a
return Trace a
tx forall (m :: * -> *) a b.
Monad m =>
m (Trace a) -> (a -> m (Trace b)) -> m (Trace b)
`bind` (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
TracedT m a -> m (WeightedT (DensityT m) a, Trace a)
runTraced forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> TracedT m b
f)
    forall (m :: * -> *) a. Monad m => a -> m a
return (WeightedT (DensityT m) b
m, Trace b
t)

instance MonadTrans TracedT where
  lift :: forall (m :: * -> *) a. Monad m => m a -> TracedT m a
lift m a
m = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,) (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m a
m) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a. Applicative f => a -> f a
pure) m a
m

instance (MonadDistribution m) => MonadDistribution (TracedT m) where
  random :: TracedT m Double
random = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((,) forall (m :: * -> *). MonadDistribution m => m Double
random forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (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 m (WeightedT (DensityT m) a, Trace a)
c) = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT (forall x. m x -> m x
f m (WeightedT (DensityT m) a, Trace a)
c)

-- | Discard the trace and supporting infrastructure.
marginal :: (Monad m) => TracedT m a -> m a
marginal :: forall (m :: * -> *) a. Monad m => TracedT m a -> m a
marginal (TracedT m (WeightedT (DensityT m) a, Trace a)
c) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Trace a -> a
output forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) m (WeightedT (DensityT m) a, Trace a)
c

-- | Freeze all traced random choices to their current values and stop tracing
-- them.
freeze :: (Monad m) => TracedT m a -> TracedT m a
freeze :: forall (m :: * -> *) a. Monad m => TracedT m a -> TracedT m a
freeze (TracedT m (WeightedT (DensityT m) a, Trace a)
c) = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ do
  (WeightedT (DensityT m) a
_, Trace a
t) <- m (WeightedT (DensityT m) a, Trace a)
c
  let x :: a
x = forall a. Trace a -> a
output Trace a
t
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall (m :: * -> *) a. Monad m => a -> m a
return a
x, forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)

-- | A single step of the Trace Metropolis-Hastings algorithm.
mhStep :: (MonadDistribution m) => TracedT m a -> TracedT m a
mhStep :: forall (m :: * -> *) a.
MonadDistribution m =>
TracedT m a -> TracedT m a
mhStep (TracedT m (WeightedT (DensityT m) a, Trace a)
c) = forall (m :: * -> *) a.
m (WeightedT (DensityT m) a, Trace a) -> TracedT m a
TracedT forall a b. (a -> b) -> a -> b
$ do
  (WeightedT (DensityT m) a
m, Trace a
t) <- m (WeightedT (DensityT m) a, Trace a)
c
  Trace a
t' <- forall (m :: * -> *) a.
MonadDistribution m =>
WeightedT (DensityT m) a -> Trace a -> m (Trace a)
mhTransFree WeightedT (DensityT m) a
m Trace a
t
  forall (m :: * -> *) a. Monad m => a -> m a
return (WeightedT (DensityT m) a
m, Trace a
t')

-- | Full run of the Trace Metropolis-Hastings algorithm with a specified
-- number of steps.
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 m (WeightedT (DensityT m) a, Trace a)
c) = do
  (WeightedT (DensityT m) a
m, Trace a
t) <- m (WeightedT (DensityT m) a, Trace a)
c
  let f :: t -> m (NonEmpty (Trace a))
f t
k
        | t
k forall a. Ord a => a -> a -> Bool
<= t
0 = forall (m :: * -> *) a. Monad m => a -> m a
return (Trace a
t forall a. a -> [a] -> NonEmpty a
:| [])
        | 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 m) a -> Trace a -> m (Trace a)
mhTransFree WeightedT (DensityT m) 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)
  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)