{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
{-# LANGUAGE StandaloneKindSignatures #-}
module Inference.Conjugate where

import           Control.Monad.Primitive        ( PrimMonad
                                                , PrimState
                                                )
import           Control.Monad.Reader           ( ReaderT
                                                , runReaderT
                                                )
import           Control.Monad.Reader.Class     ( MonadReader(ask) )
import           Control.Monad.State            ( State
                                                , StateT
                                                , evalState
                                                , evalStateT
                                                , execStateT
                                                , runStateT
                                                )
import           Control.Monad.State.Class      ( get
                                                , modify
                                                , put
                                                )
import           Control.Monad.Trans.Maybe      ( MaybeT(MaybeT)
                                                , runMaybeT
                                                )
import           Control.Monad.Writer
import           Data.Dynamic                   ( Dynamic
                                                , Typeable
                                                , fromDynamic
                                                , toDyn
                                                )
import           Data.Kind
import           Data.Maybe                     ( fromMaybe )
import           Data.MultiSet                 as MS
import qualified Data.Sequence                 as S
import           Data.Typeable                  ( Proxy(Proxy)
                                                , typeRep
                                                )
import qualified Data.Vector                   as V
import qualified Debug.Trace                   as DT
import           GHC.Float                      ( int2Double )
import           GHC.Generics
import           GHC.TypeNats
import           Lens.Micro
import           Lens.Micro.Extras
import           Lens.Micro.TH                  ( makeLenses )
import           Numeric.SpecFunctions          ( logBeta
                                                , logChoose
                                                , logFactorial
                                                , logGamma
                                                )
import           System.Random.MWC.Probability
                                         hiding ( Uniform )

----------------------------------------------------------------
-- type classes and families for distributions and conjugates --
----------------------------------------------------------------

-- | Describes a family of distributions with a fixed form.
-- For example, a 'Bernoulli' distribution is parameterized by a probability @p@
-- and produces binary samples
-- (@True@ with probability @p@, @False@ with probability @1-p@).
--
-- Its 'Distribution' instance is:
-- > instance Distribution Bernoulli where
-- >   type Params Bernoulli = Double
-- >   type Support Bernoulli = Bool
-- >   distSample _ = uncurry bernoulli
-- >   distLogP _ p True = log p
-- >   distLogP _ p False = log (1 - p)
class Distribution a where
  type Params a :: Type
  type Support a :: Type
  distSample :: (PrimMonad m) => a -> Params a -> Prob m (Support a)
  distLogP :: a -> Params a -> Support a -> Double

-- -- | Used as a type-lifted kind for conjugate distribution pairs.
-- data Conj a b = Conj a bb

-- | A type-level marker for treating a distribution as a prior.
newtype AsPrior p = AsPrior p

-- -- | A type-level marker for treating a distribution as a likelihood.
-- newtype AsLk l = AsLk l

-- | Marks two distributions as a conjugate pair of prior and likelihood.
-- The property of such a pair is that the posterior has the same form as the prior
-- (including the same 'Params' and 'Support'),
-- and that its parameters can be obtained analytically from the parameters of the prior
-- and a set of observations.
--
-- The class method 'updatePrior' returns the parameters of the posterior
-- given the prior parameters after a single observation.
class (Distribution p, Distribution l, Support p ~ Params l) => Conjugate p l where
  priorSingleton :: p -- ^ provides a singleton instance of the prior distribution
                      -- in order to make sampling from priors easier.
  updatePrior :: l -> Params p -> Support l -> Params p
  predLogP :: l -> Params p -> Support l -> Double

type family Hyper (a :: k) :: Type
type instance Hyper (AsPrior p) = Params p

type family Probs (a :: k) :: Type
type instance Probs (AsPrior p) = Support p

-- type family Value (a :: k) :: Type
-- type instance Value (Conj p l) = Support l

-- helper types for instantiating hyperparameters, parameters, and values
-- ----------------------------------------------------------------------

newtype HyperRep p = HyperRep { forall p. HyperRep p -> Hyper (AsPrior p)
runHyper :: Hyper (AsPrior p) }
deriving instance Show (Hyper (AsPrior p)) => Show (HyperRep p)

type instance Hyper (a :: (Type -> Type) -> Type) = a HyperRep

newtype ProbsRep p = ProbsRep { forall p. ProbsRep p -> Probs (AsPrior p)
runProbs :: Probs (AsPrior p)}
deriving instance Show (Probs (AsPrior p)) => Show (ProbsRep p)

type instance Probs (a :: (Type -> Type) -> Type) = a ProbsRep

-- newtype ValueRep p l = ValueRep { runValue :: Value (Conj p l) }
-- deriving instance Show (Value (Conj p l)) => Show (ValueRep p l)

-- type instance Value (a :: (Type -> Type -> Type) -> Type) = a ValueRep

-----------------------------------------------------
-- generic magic for conjugates and parameter sets --
-----------------------------------------------------

-- Jeffreys prior
-- ---------------

class Jeffreys a where
  jeffreysPrior :: Hyper a

class GJeffreys t where
  gjeffreysPrior :: forall p. t p

instance GJeffreys V1 where
  gjeffreysPrior :: forall (p :: k). V1 p
gjeffreysPrior = forall a. HasCallStack => a
undefined -- ok

instance GJeffreys U1 where
  gjeffreysPrior :: forall (p :: k). U1 p
gjeffreysPrior = forall k (p :: k). U1 p
U1

-- base case: k is a conjugate distribution
instance (Jeffreys (AsPrior p)) => GJeffreys (K1 i (HyperRep p)) where
  gjeffreysPrior :: forall (p :: k). K1 i (HyperRep p) p
gjeffreysPrior = forall k i c (p :: k). c -> K1 i c p
K1 forall a b. (a -> b) -> a -> b
$ forall p. Hyper (AsPrior p) -> HyperRep p
HyperRep forall a b. (a -> b) -> a -> b
$ forall {k} (a :: k). Jeffreys a => Hyper a
jeffreysPrior @(AsPrior p)

-- recursive case: k is another record
instance (Jeffreys k, k HyperRep ~ Hyper k) => GJeffreys (K1 i (k HyperRep)) where
  gjeffreysPrior :: forall (p :: k). K1 i (k HyperRep) p
gjeffreysPrior = forall k i c (p :: k). c -> K1 i c p
K1 forall a b. (a -> b) -> a -> b
$ forall {k} (a :: k). Jeffreys a => Hyper a
jeffreysPrior @k

instance (GJeffreys t) => GJeffreys (M1 i c (t :: Type -> Type)) where
  gjeffreysPrior :: forall p. M1 i c t p
gjeffreysPrior = forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 forall {k} (t :: k -> *) (p :: k). GJeffreys t => t p
gjeffreysPrior

instance (GJeffreys ta, GJeffreys tb) => GJeffreys (ta :*: tb) where
  gjeffreysPrior :: forall (p :: k). (:*:) ta tb p
gjeffreysPrior = forall {k} (t :: k -> *) (p :: k). GJeffreys t => t p
gjeffreysPrior @ta forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: forall {k} (t :: k -> *) (p :: k). GJeffreys t => t p
gjeffreysPrior @tb

instance (Generic (t HyperRep), GJeffreys (Rep (t HyperRep))) => Jeffreys (t :: (Type -> Type) -> Type) where
  jeffreysPrior :: Hyper t
jeffreysPrior = forall a x. Generic a => Rep a x -> a
GHC.Generics.to (forall {k} (t :: k -> *) (p :: k). GJeffreys t => t p
gjeffreysPrior @(Rep (t HyperRep)))

-- uniform prior
-- -------------

class Uniform a where
  uniformPrior :: Hyper a

class GUniform t where
  guniformPrior :: forall p. t p

instance GUniform V1 where
  guniformPrior :: forall (p :: k). V1 p
guniformPrior = forall a. HasCallStack => a
undefined -- ok

instance GUniform U1 where
  guniformPrior :: forall (p :: k). U1 p
guniformPrior = forall k (p :: k). U1 p
U1

-- base case: k is a conjugate distribution
instance (Uniform (AsPrior p)) => GUniform (K1 i (HyperRep p)) where
  guniformPrior :: forall (p :: k). K1 i (HyperRep p) p
guniformPrior = forall k i c (p :: k). c -> K1 i c p
K1 forall a b. (a -> b) -> a -> b
$ forall p. Hyper (AsPrior p) -> HyperRep p
HyperRep forall a b. (a -> b) -> a -> b
$ forall {k} (a :: k). Uniform a => Hyper a
uniformPrior @(AsPrior p)

-- recursive case: k is another record
instance (Uniform k, k HyperRep ~ Hyper k) => GUniform (K1 i (k HyperRep)) where
  guniformPrior :: forall (p :: k). K1 i (k HyperRep) p
guniformPrior = forall k i c (p :: k). c -> K1 i c p
K1 forall a b. (a -> b) -> a -> b
$ forall {k} (a :: k). Uniform a => Hyper a
uniformPrior @k

instance (GUniform t) => GUniform (M1 i c (t :: Type -> Type)) where
  guniformPrior :: forall p. M1 i c t p
guniformPrior = forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 forall {k} (t :: k -> *) (p :: k). GUniform t => t p
guniformPrior

instance (GUniform ta, GUniform tb) => GUniform (ta :*: tb) where
  guniformPrior :: forall (p :: k). (:*:) ta tb p
guniformPrior = forall {k} (t :: k -> *) (p :: k). GUniform t => t p
guniformPrior @ta forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
:*: forall {k} (t :: k -> *) (p :: k). GUniform t => t p
guniformPrior @tb

instance (Generic (t HyperRep), GUniform (Rep (t HyperRep))) => Uniform (t :: (Type -> Type) -> Type) where
  uniformPrior :: Hyper t
uniformPrior = forall a x. Generic a => Rep a x -> a
GHC.Generics.to (forall {k} (t :: k -> *) (p :: k). GUniform t => t p
guniformPrior @(Rep (t HyperRep)))

-- sampling from prior
-- ------------------

class Prior a where
  sampleProbs :: (PrimMonad m) => Hyper a -> Prob m (Probs a)

class GPrior i o where
  gsampleProbs :: forall m p. PrimMonad m => i p -> Prob m (o p)

instance GPrior V1 V1 where
  gsampleProbs :: forall (m :: * -> *) (p :: k). PrimMonad m => V1 p -> Prob m (V1 p)
gsampleProbs = forall a. HasCallStack => a
undefined -- ok

instance GPrior U1 U1 where
  gsampleProbs :: forall (m :: * -> *) (p :: k). PrimMonad m => U1 p -> Prob m (U1 p)
gsampleProbs U1 p
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall k (p :: k). U1 p
U1

-- base case: k is a conjugate distribution
instance (Prior (AsPrior p)) => GPrior (K1 i (HyperRep p)) (K1 i (ProbsRep p)) where
  gsampleProbs :: forall (m :: * -> *) (p :: k).
PrimMonad m =>
K1 i (HyperRep p) p -> Prob m (K1 i (ProbsRep p) p)
gsampleProbs (K1 (HyperRep Hyper (AsPrior p)
hyper)) =
    forall k i c (p :: k). c -> K1 i c p
K1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall p. Probs (AsPrior p) -> ProbsRep p
ProbsRep forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (a :: k) (m :: * -> *).
(Prior a, PrimMonad m) =>
Hyper a -> Prob m (Probs a)
sampleProbs @(AsPrior p) Hyper (AsPrior p)
hyper

-- recursive case: k is another record
instance (Prior k, k HyperRep ~ Hyper k, k ProbsRep ~ Probs k) =>
         GPrior (K1 i (k HyperRep)) (K1 i (k ProbsRep)) where
  gsampleProbs :: forall (m :: * -> *) (p :: k).
PrimMonad m =>
K1 i (k HyperRep) p -> Prob m (K1 i (k ProbsRep) p)
gsampleProbs (K1 k HyperRep
hyper) = forall k i c (p :: k). c -> K1 i c p
K1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (a :: k) (m :: * -> *).
(Prior a, PrimMonad m) =>
Hyper a -> Prob m (Probs a)
sampleProbs @k k HyperRep
hyper

instance (GPrior ti to) => GPrior (M1 i c ti) (M1 i' c' to) where
  gsampleProbs :: forall (m :: * -> *) (p :: k).
PrimMonad m =>
M1 i c ti p -> Prob m (M1 i' c' to p)
gsampleProbs (M1 ti p
x) = forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (i :: k -> *) (o :: k -> *) (m :: * -> *) (p :: k).
(GPrior i o, PrimMonad m) =>
i p -> Prob m (o p)
gsampleProbs ti p
x

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :*: ib) (oa :*: ob) where
  gsampleProbs :: forall (m :: * -> *) (p :: k).
PrimMonad m =>
(:*:) ia ib p -> Prob m ((:*:) oa ob p)
gsampleProbs (ia p
a :*: ib p
b) = forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
(:*:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (i :: k -> *) (o :: k -> *) (m :: * -> *) (p :: k).
(GPrior i o, PrimMonad m) =>
i p -> Prob m (o p)
gsampleProbs ia p
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall {k} (i :: k -> *) (o :: k -> *) (m :: * -> *) (p :: k).
(GPrior i o, PrimMonad m) =>
i p -> Prob m (o p)
gsampleProbs ib p
b

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :+: ib) (oa :+: ob) where
  gsampleProbs :: forall (m :: * -> *) (p :: k).
PrimMonad m =>
(:+:) ia ib p -> Prob m ((:+:) oa ob p)
gsampleProbs (L1 ia p
a) = forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (i :: k -> *) (o :: k -> *) (m :: * -> *) (p :: k).
(GPrior i o, PrimMonad m) =>
i p -> Prob m (o p)
gsampleProbs ia p
a
  gsampleProbs (R1 ib p
b) = forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (i :: k -> *) (o :: k -> *) (m :: * -> *) (p :: k).
(GPrior i o, PrimMonad m) =>
i p -> Prob m (o p)
gsampleProbs ib p
b

instance ( Generic (a HyperRep)
         , Generic (a ProbsRep)
         , GPrior (Rep (a HyperRep)) (Rep (a ProbsRep))
         ) => Prior (a :: (Type -> Type) -> Type) where
  sampleProbs :: forall (m :: * -> *). PrimMonad m => Hyper a -> Prob m (Probs a)
sampleProbs Hyper a
hyper = forall a x. Generic a => Rep a x -> a
GHC.Generics.to forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (i :: k -> *) (o :: k -> *) (m :: * -> *) (p :: k).
(GPrior i o, PrimMonad m) =>
i p -> Prob m (o p)
gsampleProbs (forall a x. Generic a => a -> Rep a x
from Hyper a
hyper)

-----------------------
-- likelihood monads --
-----------------------

type Accessor r p = forall f . Lens' (r f) (f p)

class Monad m => RandomInterpreter m r | m -> r where
  type SampleCtx m a :: Constraint
  sampleValue :: (Conjugate p l, SampleCtx m l) => String -> l -> Accessor r p -> m (Support l)
  sampleConst :: (Distribution d, SampleCtx m d) => String -> d -> Params d -> m (Support d)
  permutationPlate :: (Ord a) => Int -> m a -> m [a]

newtype Trace (r :: (Type -> Type) -> Type)  = Trace { forall (r :: (* -> *) -> *). Trace r -> Seq Dynamic
runTrace :: S.Seq Dynamic }
  deriving (Int -> Trace r -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (r :: (* -> *) -> *). Int -> Trace r -> ShowS
forall (r :: (* -> *) -> *). [Trace r] -> ShowS
forall (r :: (* -> *) -> *). Trace r -> String
showList :: [Trace r] -> ShowS
$cshowList :: forall (r :: (* -> *) -> *). [Trace r] -> ShowS
show :: Trace r -> String
$cshow :: forall (r :: (* -> *) -> *). Trace r -> String
showsPrec :: Int -> Trace r -> ShowS
$cshowsPrec :: forall (r :: (* -> *) -> *). Int -> Trace r -> ShowS
Show)

observeValue
  :: (Conjugate p l, Typeable (Support l), Monad m)
  => String
  -> l
  -> Accessor r p
  -> Support l
  -> StateT (Trace r) m ()
observeValue :: forall p l (m :: * -> *) (r :: (* -> *) -> *).
(Conjugate p l, Typeable (Support l), Monad m) =>
String -> l -> Accessor r p -> Support l -> StateT (Trace r) m ()
observeValue String
_ l
_ Accessor r p
_ Support l
val = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \(Trace Seq Dynamic
st) -> forall (r :: (* -> *) -> *). Seq Dynamic -> Trace r
Trace forall a b. (a -> b) -> a -> b
$ Seq Dynamic
st forall a. Seq a -> a -> Seq a
S.|> forall a. Typeable a => a -> Dynamic
toDyn Support l
val

observeConst
  :: (Distribution d, Typeable (Support d), Monad m)
  => String
  -> d
  -> Params d
  -> Support d
  -> StateT (Trace r) m ()
observeConst :: forall d (m :: * -> *) (r :: (* -> *) -> *).
(Distribution d, Typeable (Support d), Monad m) =>
String -> d -> Params d -> Support d -> StateT (Trace r) m ()
observeConst String
_ d
_ Params d
_ Support d
val = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \(Trace Seq Dynamic
st) -> forall (r :: (* -> *) -> *). Seq Dynamic -> Trace r
Trace forall a b. (a -> b) -> a -> b
$ Seq Dynamic
st forall a. Seq a -> a -> Seq a
S.|> forall a. Typeable a => a -> Dynamic
toDyn Support d
val

takeTrace :: Typeable a => Trace r -> Maybe (a, Trace r)
takeTrace :: forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace (Trace Seq Dynamic
t) = do
  (Dynamic
valDyn, Seq Dynamic
rest) <- case forall a. Seq a -> ViewL a
S.viewl Seq Dynamic
t of
    ViewL Dynamic
S.EmptyL         -> forall a. Maybe a
Nothing
    Dynamic
valDyn S.:< Seq Dynamic
rest -> forall a. a -> Maybe a
Just (Dynamic
valDyn, Seq Dynamic
rest)
  a
val <- forall a. Typeable a => Dynamic -> Maybe a
fromDynamic Dynamic
valDyn
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
val, forall (r :: (* -> *) -> *). Seq Dynamic -> Trace r
Trace Seq Dynamic
rest)

-- just sample
-- -----------

newtype SampleI m r a = SampleI (ReaderT (r ProbsRep) (Prob m) a)
  deriving (forall a b. a -> SampleI m r b -> SampleI m r a
forall a b. (a -> b) -> SampleI m r a -> SampleI m r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
a -> SampleI m r b -> SampleI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
(a -> b) -> SampleI m r a -> SampleI m r b
<$ :: forall a b. a -> SampleI m r b -> SampleI m r a
$c<$ :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
a -> SampleI m r b -> SampleI m r a
fmap :: forall a b. (a -> b) -> SampleI m r a -> SampleI m r b
$cfmap :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
(a -> b) -> SampleI m r a -> SampleI m r b
Functor, forall a. a -> SampleI m r a
forall a b. SampleI m r a -> SampleI m r b -> SampleI m r a
forall a b. SampleI m r a -> SampleI m r b -> SampleI m r b
forall a b. SampleI m r (a -> b) -> SampleI m r a -> SampleI m r b
forall a b c.
(a -> b -> c) -> SampleI m r a -> SampleI m r b -> SampleI m r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall {m :: * -> *} {r :: (* -> *) -> *}.
Monad m =>
Functor (SampleI m r)
forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> SampleI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> SampleI m r b -> SampleI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> SampleI m r b -> SampleI m r b
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r (a -> b) -> SampleI m r a -> SampleI m r b
forall (m :: * -> *) (r :: (* -> *) -> *) a b c.
Monad m =>
(a -> b -> c) -> SampleI m r a -> SampleI m r b -> SampleI m r c
<* :: forall a b. SampleI m r a -> SampleI m r b -> SampleI m r a
$c<* :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> SampleI m r b -> SampleI m r a
*> :: forall a b. SampleI m r a -> SampleI m r b -> SampleI m r b
$c*> :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> SampleI m r b -> SampleI m r b
liftA2 :: forall a b c.
(a -> b -> c) -> SampleI m r a -> SampleI m r b -> SampleI m r c
$cliftA2 :: forall (m :: * -> *) (r :: (* -> *) -> *) a b c.
Monad m =>
(a -> b -> c) -> SampleI m r a -> SampleI m r b -> SampleI m r c
<*> :: forall a b. SampleI m r (a -> b) -> SampleI m r a -> SampleI m r b
$c<*> :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r (a -> b) -> SampleI m r a -> SampleI m r b
pure :: forall a. a -> SampleI m r a
$cpure :: forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> SampleI m r a
Applicative, forall a. a -> SampleI m r a
forall a b. SampleI m r a -> SampleI m r b -> SampleI m r b
forall a b. SampleI m r a -> (a -> SampleI m r b) -> SampleI m r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (m :: * -> *) (r :: (* -> *) -> *).
Monad m =>
Applicative (SampleI m r)
forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> SampleI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> SampleI m r b -> SampleI m r b
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> (a -> SampleI m r b) -> SampleI m r b
return :: forall a. a -> SampleI m r a
$creturn :: forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> SampleI m r a
>> :: forall a b. SampleI m r a -> SampleI m r b -> SampleI m r b
$c>> :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> SampleI m r b -> SampleI m r b
>>= :: forall a b. SampleI m r a -> (a -> SampleI m r b) -> SampleI m r b
$c>>= :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
SampleI m r a -> (a -> SampleI m r b) -> SampleI m r b
Monad)

instance (PrimMonad m) => RandomInterpreter (SampleI m r) r where
  type SampleCtx (SampleI m r) a = ()
  sampleValue
    :: forall p l
     . (Conjugate p l)
    => String
    -> l
    -> Accessor r p
    -> SampleI m r (Support l)
  sampleValue :: forall p l.
Conjugate p l =>
String -> l -> Accessor r p -> SampleI m r (Support l)
sampleValue String
_ l
lk Accessor r p
getProbs = forall (m :: * -> *) (r :: (* -> *) -> *) a.
ReaderT (r ProbsRep) (Prob m) a -> SampleI m r a
SampleI forall a b. (a -> b) -> a -> b
$ do
    r ProbsRep
probs <- forall r (m :: * -> *). MonadReader r m => m r
ask
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(Distribution a, PrimMonad m) =>
a -> Params a -> Prob m (Support a)
distSample l
lk forall a b. (a -> b) -> a -> b
$ forall p. ProbsRep p -> Probs (AsPrior p)
runProbs forall a b. (a -> b) -> a -> b
$ forall a s. Getting a s a -> s -> a
view Accessor r p
getProbs r ProbsRep
probs
  sampleConst
    :: forall d
     . (Distribution d)
    => String
    -> d
    -> Params d
    -> SampleI m r (Support d)
  sampleConst :: forall d.
Distribution d =>
String -> d -> Params d -> SampleI m r (Support d)
sampleConst String
_ d
dist Params d
params = forall (m :: * -> *) (r :: (* -> *) -> *) a.
ReaderT (r ProbsRep) (Prob m) a -> SampleI m r a
SampleI forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(Distribution a, PrimMonad m) =>
a -> Params a -> Prob m (Support a)
distSample d
dist Params d
params
  permutationPlate :: forall a. Ord a => Int -> SampleI m r a -> SampleI m r [a]
permutationPlate = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM

sampleResult :: p ProbsRep -> SampleI m p a -> Gen (PrimState m) -> m a
sampleResult :: forall (p :: (* -> *) -> *) (m :: * -> *) a.
p ProbsRep -> SampleI m p a -> Gen (PrimState m) -> m a
sampleResult p ProbsRep
probs (SampleI ReaderT (p ProbsRep) (Prob m) a
a) = forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (p ProbsRep) (Prob m) a
a p ProbsRep
probs)

-- sample and trace the execution process
-- --------------------------------------

newtype TraceI m r a = TraceI (ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a)
  deriving (forall a b. a -> TraceI m r b -> TraceI m r a
forall a b. (a -> b) -> TraceI m r a -> TraceI m r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
a -> TraceI m r b -> TraceI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
(a -> b) -> TraceI m r a -> TraceI m r b
<$ :: forall a b. a -> TraceI m r b -> TraceI m r a
$c<$ :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
a -> TraceI m r b -> TraceI m r a
fmap :: forall a b. (a -> b) -> TraceI m r a -> TraceI m r b
$cfmap :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Functor m =>
(a -> b) -> TraceI m r a -> TraceI m r b
Functor, forall a. a -> TraceI m r a
forall a b. TraceI m r a -> TraceI m r b -> TraceI m r a
forall a b. TraceI m r a -> TraceI m r b -> TraceI m r b
forall a b. TraceI m r (a -> b) -> TraceI m r a -> TraceI m r b
forall a b c.
(a -> b -> c) -> TraceI m r a -> TraceI m r b -> TraceI m r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall {m :: * -> *} {r :: (* -> *) -> *}.
Monad m =>
Functor (TraceI m r)
forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> TraceI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> TraceI m r b -> TraceI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> TraceI m r b -> TraceI m r b
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r (a -> b) -> TraceI m r a -> TraceI m r b
forall (m :: * -> *) (r :: (* -> *) -> *) a b c.
Monad m =>
(a -> b -> c) -> TraceI m r a -> TraceI m r b -> TraceI m r c
<* :: forall a b. TraceI m r a -> TraceI m r b -> TraceI m r a
$c<* :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> TraceI m r b -> TraceI m r a
*> :: forall a b. TraceI m r a -> TraceI m r b -> TraceI m r b
$c*> :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> TraceI m r b -> TraceI m r b
liftA2 :: forall a b c.
(a -> b -> c) -> TraceI m r a -> TraceI m r b -> TraceI m r c
$cliftA2 :: forall (m :: * -> *) (r :: (* -> *) -> *) a b c.
Monad m =>
(a -> b -> c) -> TraceI m r a -> TraceI m r b -> TraceI m r c
<*> :: forall a b. TraceI m r (a -> b) -> TraceI m r a -> TraceI m r b
$c<*> :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r (a -> b) -> TraceI m r a -> TraceI m r b
pure :: forall a. a -> TraceI m r a
$cpure :: forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> TraceI m r a
Applicative, forall a. a -> TraceI m r a
forall a b. TraceI m r a -> TraceI m r b -> TraceI m r b
forall a b. TraceI m r a -> (a -> TraceI m r b) -> TraceI m r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (m :: * -> *) (r :: (* -> *) -> *).
Monad m =>
Applicative (TraceI m r)
forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> TraceI m r a
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> TraceI m r b -> TraceI m r b
forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> (a -> TraceI m r b) -> TraceI m r b
return :: forall a. a -> TraceI m r a
$creturn :: forall (m :: * -> *) (r :: (* -> *) -> *) a.
Monad m =>
a -> TraceI m r a
>> :: forall a b. TraceI m r a -> TraceI m r b -> TraceI m r b
$c>> :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> TraceI m r b -> TraceI m r b
>>= :: forall a b. TraceI m r a -> (a -> TraceI m r b) -> TraceI m r b
$c>>= :: forall (m :: * -> *) (r :: (* -> *) -> *) a b.
Monad m =>
TraceI m r a -> (a -> TraceI m r b) -> TraceI m r b
Monad)

instance (PrimMonad m) => RandomInterpreter (TraceI m r) r where
  type SampleCtx (TraceI m r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceI m r (Support l)
  sampleValue :: forall p l.
(Conjugate p l, Typeable (Support l)) =>
String -> l -> Accessor r p -> TraceI m r (Support l)
sampleValue String
_ l
lk Accessor r p
getProbs = forall (m :: * -> *) (r :: (* -> *) -> *) a.
ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a -> TraceI m r a
TraceI forall a b. (a -> b) -> a -> b
$ do
    r ProbsRep
probs <- forall r (m :: * -> *). MonadReader r m => m r
ask
    Support l
val   <- 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 forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(Distribution a, PrimMonad m) =>
a -> Params a -> Prob m (Support a)
distSample l
lk forall a b. (a -> b) -> a -> b
$ forall p. ProbsRep p -> Probs (AsPrior p)
runProbs forall a b. (a -> b) -> a -> b
$ forall a s. Getting a s a -> s -> a
view Accessor r p
getProbs r ProbsRep
probs
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \(Trace Seq Dynamic
obs) -> forall (r :: (* -> *) -> *). Seq Dynamic -> Trace r
Trace forall a b. (a -> b) -> a -> b
$ Seq Dynamic
obs forall a. Seq a -> a -> Seq a
S.|> forall a. Typeable a => a -> Dynamic
toDyn Support l
val
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support l
val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> TraceI m r (Support d)
  sampleConst :: forall d.
(Distribution d, Typeable (Support d)) =>
String -> d -> Params d -> TraceI m r (Support d)
sampleConst String
_ d
dist Params d
params = forall (m :: * -> *) (r :: (* -> *) -> *) a.
ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a -> TraceI m r a
TraceI forall a b. (a -> b) -> a -> b
$ do
    Support d
val <- 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 forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(Distribution a, PrimMonad m) =>
a -> Params a -> Prob m (Support a)
distSample d
dist Params d
params
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \(Trace Seq Dynamic
obs) -> forall (r :: (* -> *) -> *). Seq Dynamic -> Trace r
Trace forall a b. (a -> b) -> a -> b
$ Seq Dynamic
obs forall a. Seq a -> a -> Seq a
S.|> forall a. Typeable a => a -> Dynamic
toDyn Support d
val
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support d
val
  permutationPlate :: forall a. Ord a => Int -> TraceI m r a -> TraceI m r [a]
permutationPlate = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM

sampleTrace
  :: r ProbsRep -> TraceI m r a -> Gen (PrimState m) -> m (a, Trace r)
sampleTrace :: forall (r :: (* -> *) -> *) (m :: * -> *) a.
r ProbsRep -> TraceI m r a -> Gen (PrimState m) -> m (a, Trace r)
sampleTrace r ProbsRep
probs (TraceI ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a
a) = do
  let st :: StateT (Trace r) (Prob m) a
st = forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a
a r ProbsRep
probs
      pr :: Prob m (a, Trace r)
pr = forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT (Trace r) (Prob m) a
st (forall (r :: (* -> *) -> *). Seq Dynamic -> Trace r
Trace forall a. Monoid a => a
mempty)
  forall (m :: * -> *) a. Prob m a -> Gen (PrimState m) -> m a
sample Prob m (a, Trace r)
pr

-- evaluate the probability of a trace
-- -----------------------------------

newtype EvalTraceI r a = EvalTraceI (ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a)
  deriving (forall a b. a -> EvalTraceI r b -> EvalTraceI r a
forall a b. (a -> b) -> EvalTraceI r a -> EvalTraceI r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (r :: (* -> *) -> *) a b.
a -> EvalTraceI r b -> EvalTraceI r a
forall (r :: (* -> *) -> *) a b.
(a -> b) -> EvalTraceI r a -> EvalTraceI r b
<$ :: forall a b. a -> EvalTraceI r b -> EvalTraceI r a
$c<$ :: forall (r :: (* -> *) -> *) a b.
a -> EvalTraceI r b -> EvalTraceI r a
fmap :: forall a b. (a -> b) -> EvalTraceI r a -> EvalTraceI r b
$cfmap :: forall (r :: (* -> *) -> *) a b.
(a -> b) -> EvalTraceI r a -> EvalTraceI r b
Functor, forall a. a -> EvalTraceI r a
forall a b. EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r a
forall a b. EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
forall a b.
EvalTraceI r (a -> b) -> EvalTraceI r a -> EvalTraceI r b
forall a b c.
(a -> b -> c) -> EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (r :: (* -> *) -> *). Functor (EvalTraceI r)
forall (r :: (* -> *) -> *) a. a -> EvalTraceI r a
forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r a
forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
forall (r :: (* -> *) -> *) a b.
EvalTraceI r (a -> b) -> EvalTraceI r a -> EvalTraceI r b
forall (r :: (* -> *) -> *) a b c.
(a -> b -> c) -> EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r c
<* :: forall a b. EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r a
$c<* :: forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r a
*> :: forall a b. EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
$c*> :: forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
liftA2 :: forall a b c.
(a -> b -> c) -> EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r c
$cliftA2 :: forall (r :: (* -> *) -> *) a b c.
(a -> b -> c) -> EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r c
<*> :: forall a b.
EvalTraceI r (a -> b) -> EvalTraceI r a -> EvalTraceI r b
$c<*> :: forall (r :: (* -> *) -> *) a b.
EvalTraceI r (a -> b) -> EvalTraceI r a -> EvalTraceI r b
pure :: forall a. a -> EvalTraceI r a
$cpure :: forall (r :: (* -> *) -> *) a. a -> EvalTraceI r a
Applicative, forall a. a -> EvalTraceI r a
forall a b. EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
forall a b.
EvalTraceI r a -> (a -> EvalTraceI r b) -> EvalTraceI r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (r :: (* -> *) -> *). Applicative (EvalTraceI r)
forall (r :: (* -> *) -> *) a. a -> EvalTraceI r a
forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> (a -> EvalTraceI r b) -> EvalTraceI r b
return :: forall a. a -> EvalTraceI r a
$creturn :: forall (r :: (* -> *) -> *) a. a -> EvalTraceI r a
>> :: forall a b. EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
$c>> :: forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> EvalTraceI r b -> EvalTraceI r b
>>= :: forall a b.
EvalTraceI r a -> (a -> EvalTraceI r b) -> EvalTraceI r b
$c>>= :: forall (r :: (* -> *) -> *) a b.
EvalTraceI r a -> (a -> EvalTraceI r b) -> EvalTraceI r b
Monad)

instance RandomInterpreter (EvalTraceI r) r where
  type SampleCtx (EvalTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalTraceI r (Support l)
  sampleValue :: forall p l.
(Conjugate p l, Typeable (Support l)) =>
String -> l -> Accessor r p -> EvalTraceI r (Support l)
sampleValue String
_ l
lk Accessor r p
getProbs = forall (r :: (* -> *) -> *) a.
ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a
-> EvalTraceI r a
EvalTraceI forall a b. (a -> b) -> a -> b
$ do
    r ProbsRep
probs              <- forall r (m :: * -> *). MonadReader r m => m r
ask
    (Trace r
trace, Double
totalLogP) <- forall s (m :: * -> *). MonadState s m => m s
get
    (Support l
val  , Trace r
trace'   ) <- 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 forall a b. (a -> b) -> a -> b
$ forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace
    let logP :: Double
logP = forall a. Distribution a => a -> Params a -> Support a -> Double
distLogP l
lk (forall p. ProbsRep p -> Probs (AsPrior p)
runProbs forall a b. (a -> b) -> a -> b
$ forall a s. Getting a s a -> s -> a
view Accessor r p
getProbs r ProbsRep
probs) Support l
val
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', Double
totalLogP forall a. Num a => a -> a -> a
+ Double
logP)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support l
val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalTraceI r (Support d)
  sampleConst :: forall d.
(Distribution d, Typeable (Support d)) =>
String -> d -> Params d -> EvalTraceI r (Support d)
sampleConst String
_ d
dist Params d
params = forall (r :: (* -> *) -> *) a.
ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a
-> EvalTraceI r a
EvalTraceI forall a b. (a -> b) -> a -> b
$ do
    (Trace r
trace, Double
totalLogP) <- forall s (m :: * -> *). MonadState s m => m s
get
    (Support d
val  , Trace r
trace'   ) <- 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 forall a b. (a -> b) -> a -> b
$ forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace
    let logP :: Double
logP = forall a. Distribution a => a -> Params a -> Support a -> Double
distLogP d
dist Params d
params Support d
val
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', Double
totalLogP forall a. Num a => a -> a -> a
+ Double
logP)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support d
val
  permutationPlate :: forall a. Ord a => Int -> EvalTraceI r a -> EvalTraceI r [a]
permutationPlate Int
n EvalTraceI r a
submodel = forall (r :: (* -> *) -> *) a.
ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a
-> EvalTraceI r a
EvalTraceI forall a b. (a -> b) -> a -> b
$ do
    r ProbsRep
probs                     <- forall r (m :: * -> *). MonadReader r m => m r
ask
    (Trace r
trace  , Double
totalLogP     ) <- forall s (m :: * -> *). MonadState s m => m s
get
    ([a]
results, (Trace r
trace', Double
logP)) <- 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 forall a b. (a -> b) -> a -> b
$ forall (r :: (* -> *) -> *) a.
r ProbsRep
-> Trace r -> EvalTraceI r a -> Maybe (a, (Trace r, Double))
runTraceLogP
      r ProbsRep
probs
      Trace r
trace
      (forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n EvalTraceI r a
submodel)
    let unique :: MultiSet a
unique = forall a. Ord a => [a] -> MultiSet a
MS.fromList [a]
results
        permutations :: Double
permutations =
          forall a. Integral a => a -> Double
logFactorial Int
n forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a. Integral a => a -> Double
logFactorial forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. MultiSet a -> [(a, Int)]
MS.toOccurList MultiSet a
unique)
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', Double
totalLogP forall a. Num a => a -> a -> a
+ Double
logP forall a. Num a => a -> a -> a
+ Double
permutations)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
results

runTraceLogP
  :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, (Trace r, Double))
runTraceLogP :: forall (r :: (* -> *) -> *) a.
r ProbsRep
-> Trace r -> EvalTraceI r a -> Maybe (a, (Trace r, Double))
runTraceLogP r ProbsRep
probs Trace r
trace (EvalTraceI ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a
model) = do
  forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a
model r ProbsRep
probs) (Trace r
trace, Double
0)

evalTraceLogP :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, Double)
evalTraceLogP :: forall (r :: (* -> *) -> *) a.
r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, Double)
evalTraceLogP r ProbsRep
probs Trace r
trace EvalTraceI r a
model = do
  (a
val, (Trace r
_trace, Double
logp)) <- forall (r :: (* -> *) -> *) a.
r ProbsRep
-> Trace r -> EvalTraceI r a -> Maybe (a, (Trace r, Double))
runTraceLogP r ProbsRep
probs Trace r
trace EvalTraceI r a
model
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
val, Double
logp)


-- evaluate the predictive probability of a trace
-- -----------------------------------

newtype EvalPredTraceI r a = EvalPredTraceI (ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a)
  deriving (forall a b. a -> EvalPredTraceI r b -> EvalPredTraceI r a
forall a b. (a -> b) -> EvalPredTraceI r a -> EvalPredTraceI r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (r :: (* -> *) -> *) a b.
a -> EvalPredTraceI r b -> EvalPredTraceI r a
forall (r :: (* -> *) -> *) a b.
(a -> b) -> EvalPredTraceI r a -> EvalPredTraceI r b
<$ :: forall a b. a -> EvalPredTraceI r b -> EvalPredTraceI r a
$c<$ :: forall (r :: (* -> *) -> *) a b.
a -> EvalPredTraceI r b -> EvalPredTraceI r a
fmap :: forall a b. (a -> b) -> EvalPredTraceI r a -> EvalPredTraceI r b
$cfmap :: forall (r :: (* -> *) -> *) a b.
(a -> b) -> EvalPredTraceI r a -> EvalPredTraceI r b
Functor, forall a. a -> EvalPredTraceI r a
forall a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r a
forall a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
forall a b.
EvalPredTraceI r (a -> b)
-> EvalPredTraceI r a -> EvalPredTraceI r b
forall a b c.
(a -> b -> c)
-> EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (r :: (* -> *) -> *). Functor (EvalPredTraceI r)
forall (r :: (* -> *) -> *) a. a -> EvalPredTraceI r a
forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r a
forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r (a -> b)
-> EvalPredTraceI r a -> EvalPredTraceI r b
forall (r :: (* -> *) -> *) a b c.
(a -> b -> c)
-> EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r c
<* :: forall a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r a
$c<* :: forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r a
*> :: forall a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
$c*> :: forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
liftA2 :: forall a b c.
(a -> b -> c)
-> EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r c
$cliftA2 :: forall (r :: (* -> *) -> *) a b c.
(a -> b -> c)
-> EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r c
<*> :: forall a b.
EvalPredTraceI r (a -> b)
-> EvalPredTraceI r a -> EvalPredTraceI r b
$c<*> :: forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r (a -> b)
-> EvalPredTraceI r a -> EvalPredTraceI r b
pure :: forall a. a -> EvalPredTraceI r a
$cpure :: forall (r :: (* -> *) -> *) a. a -> EvalPredTraceI r a
Applicative, forall a. a -> EvalPredTraceI r a
forall a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
forall a b.
EvalPredTraceI r a
-> (a -> EvalPredTraceI r b) -> EvalPredTraceI r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (r :: (* -> *) -> *). Applicative (EvalPredTraceI r)
forall (r :: (* -> *) -> *) a. a -> EvalPredTraceI r a
forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a
-> (a -> EvalPredTraceI r b) -> EvalPredTraceI r b
return :: forall a. a -> EvalPredTraceI r a
$creturn :: forall (r :: (* -> *) -> *) a. a -> EvalPredTraceI r a
>> :: forall a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
$c>> :: forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a -> EvalPredTraceI r b -> EvalPredTraceI r b
>>= :: forall a b.
EvalPredTraceI r a
-> (a -> EvalPredTraceI r b) -> EvalPredTraceI r b
$c>>= :: forall (r :: (* -> *) -> *) a b.
EvalPredTraceI r a
-> (a -> EvalPredTraceI r b) -> EvalPredTraceI r b
Monad)

instance RandomInterpreter (EvalPredTraceI r) r where
  type SampleCtx (EvalPredTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalPredTraceI r (Support l)
  sampleValue :: forall p l.
(Conjugate p l, Typeable (Support l)) =>
String -> l -> Accessor r p -> EvalPredTraceI r (Support l)
sampleValue String
_ l
lk Accessor r p
getHyper = forall (r :: (* -> *) -> *) a.
ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a
-> EvalPredTraceI r a
EvalPredTraceI forall a b. (a -> b) -> a -> b
$ do
    r HyperRep
hyper              <- forall r (m :: * -> *). MonadReader r m => m r
ask
    (Trace r
trace, Double
totalLogP) <- forall s (m :: * -> *). MonadState s m => m s
get
    (Support l
val  , Trace r
trace'   ) <- 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 forall a b. (a -> b) -> a -> b
$ forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace
    let logP :: Double
logP = forall p l. Conjugate p l => l -> Params p -> Support l -> Double
predLogP @p l
lk (forall p. HyperRep p -> Hyper (AsPrior p)
runHyper forall a b. (a -> b) -> a -> b
$ forall a s. Getting a s a -> s -> a
view Accessor r p
getHyper r HyperRep
hyper) Support l
val
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', Double
totalLogP forall a. Num a => a -> a -> a
+ Double
logP)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support l
val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalPredTraceI r (Support d)
  sampleConst :: forall d.
(Distribution d, Typeable (Support d)) =>
String -> d -> Params d -> EvalPredTraceI r (Support d)
sampleConst String
_ d
dist Params d
params = forall (r :: (* -> *) -> *) a.
ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a
-> EvalPredTraceI r a
EvalPredTraceI forall a b. (a -> b) -> a -> b
$ do
    (Trace r
trace, Double
totalLogP) <- forall s (m :: * -> *). MonadState s m => m s
get
    (Support d
val  , Trace r
trace'   ) <- 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 forall a b. (a -> b) -> a -> b
$ forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace
    let logP :: Double
logP = forall a. Distribution a => a -> Params a -> Support a -> Double
distLogP d
dist Params d
params Support d
val
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', Double
totalLogP forall a. Num a => a -> a -> a
+ Double
logP)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support d
val
  permutationPlate :: forall a.
Ord a =>
Int -> EvalPredTraceI r a -> EvalPredTraceI r [a]
permutationPlate Int
n EvalPredTraceI r a
submodel = forall (r :: (* -> *) -> *) a.
ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a
-> EvalPredTraceI r a
EvalPredTraceI forall a b. (a -> b) -> a -> b
$ do
    r HyperRep
probs                     <- forall r (m :: * -> *). MonadReader r m => m r
ask
    (Trace r
trace  , Double
totalLogP     ) <- forall s (m :: * -> *). MonadState s m => m s
get
    ([a]
results, (Trace r
trace', Double
logP)) <- 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 forall a b. (a -> b) -> a -> b
$ forall (r :: (* -> *) -> *) a.
r HyperRep
-> Trace r -> EvalPredTraceI r a -> Maybe (a, (Trace r, Double))
runTracePredLogP
      r HyperRep
probs
      Trace r
trace
      (forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n EvalPredTraceI r a
submodel)
    let unique :: MultiSet a
unique = forall a. Ord a => [a] -> MultiSet a
MS.fromList [a]
results
        permutations :: Double
permutations =
          forall a. Integral a => a -> Double
logFactorial Int
n forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a. Integral a => a -> Double
logFactorial forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. MultiSet a -> [(a, Int)]
MS.toOccurList MultiSet a
unique)
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', Double
totalLogP forall a. Num a => a -> a -> a
+ Double
logP forall a. Num a => a -> a -> a
+ Double
permutations)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
results

runTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, (Trace r, Double))
runTracePredLogP :: forall (r :: (* -> *) -> *) a.
r HyperRep
-> Trace r -> EvalPredTraceI r a -> Maybe (a, (Trace r, Double))
runTracePredLogP r HyperRep
hyper Trace r
trace (EvalPredTraceI ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a
model) = do
  forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a
model r HyperRep
hyper) (Trace r
trace, Double
0)

evalTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, Double)
evalTracePredLogP :: forall (r :: (* -> *) -> *) a.
r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, Double)
evalTracePredLogP r HyperRep
hyper Trace r
trace EvalPredTraceI r a
model = do
  (a
val, (Trace r
_trace, Double
logp)) <- forall (r :: (* -> *) -> *) a.
r HyperRep
-> Trace r -> EvalPredTraceI r a -> Maybe (a, (Trace r, Double))
runTracePredLogP r HyperRep
hyper Trace r
trace EvalPredTraceI r a
model
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
val, Double
logp)


-- update priors
-- -------------

newtype UpdatePriorsI r a = UpdatePriorsI (StateT (Trace r, r HyperRep) Maybe a)
  deriving (forall a b. a -> UpdatePriorsI r b -> UpdatePriorsI r a
forall a b. (a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (r :: (* -> *) -> *) a b.
a -> UpdatePriorsI r b -> UpdatePriorsI r a
forall (r :: (* -> *) -> *) a b.
(a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
<$ :: forall a b. a -> UpdatePriorsI r b -> UpdatePriorsI r a
$c<$ :: forall (r :: (* -> *) -> *) a b.
a -> UpdatePriorsI r b -> UpdatePriorsI r a
fmap :: forall a b. (a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
$cfmap :: forall (r :: (* -> *) -> *) a b.
(a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
Functor, forall a. a -> UpdatePriorsI r a
forall a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r a
forall a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
forall a b.
UpdatePriorsI r (a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
forall a b c.
(a -> b -> c)
-> UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (r :: (* -> *) -> *). Functor (UpdatePriorsI r)
forall (r :: (* -> *) -> *) a. a -> UpdatePriorsI r a
forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r a
forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r (a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
forall (r :: (* -> *) -> *) a b c.
(a -> b -> c)
-> UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r c
<* :: forall a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r a
$c<* :: forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r a
*> :: forall a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
$c*> :: forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
liftA2 :: forall a b c.
(a -> b -> c)
-> UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r c
$cliftA2 :: forall (r :: (* -> *) -> *) a b c.
(a -> b -> c)
-> UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r c
<*> :: forall a b.
UpdatePriorsI r (a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
$c<*> :: forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r (a -> b) -> UpdatePriorsI r a -> UpdatePriorsI r b
pure :: forall a. a -> UpdatePriorsI r a
$cpure :: forall (r :: (* -> *) -> *) a. a -> UpdatePriorsI r a
Applicative, forall a. a -> UpdatePriorsI r a
forall a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
forall a b.
UpdatePriorsI r a -> (a -> UpdatePriorsI r b) -> UpdatePriorsI r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (r :: (* -> *) -> *). Applicative (UpdatePriorsI r)
forall (r :: (* -> *) -> *) a. a -> UpdatePriorsI r a
forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> (a -> UpdatePriorsI r b) -> UpdatePriorsI r b
return :: forall a. a -> UpdatePriorsI r a
$creturn :: forall (r :: (* -> *) -> *) a. a -> UpdatePriorsI r a
>> :: forall a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
$c>> :: forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> UpdatePriorsI r b -> UpdatePriorsI r b
>>= :: forall a b.
UpdatePriorsI r a -> (a -> UpdatePriorsI r b) -> UpdatePriorsI r b
$c>>= :: forall (r :: (* -> *) -> *) a b.
UpdatePriorsI r a -> (a -> UpdatePriorsI r b) -> UpdatePriorsI r b
Monad)

instance RandomInterpreter (UpdatePriorsI r) r where
  type SampleCtx (UpdatePriorsI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> UpdatePriorsI r (Support l)
  sampleValue :: forall p l.
(Conjugate p l, Typeable (Support l)) =>
String -> l -> Accessor r p -> UpdatePriorsI r (Support l)
sampleValue String
_ l
lk Accessor r p
accessor = forall (r :: (* -> *) -> *) a.
StateT (Trace r, r HyperRep) Maybe a -> UpdatePriorsI r a
UpdatePriorsI forall a b. (a -> b) -> a -> b
$ do
    (Trace r
trace, r HyperRep
priors) <- forall s (m :: * -> *). MonadState s m => m s
get
    (Support l
val  , Trace r
trace') <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace
    let priors' :: r HyperRep
        priors' :: r HyperRep
priors' = forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over
          Accessor r p
accessor
          (\(HyperRep Hyper (AsPrior p)
pr) -> forall p. Hyper (AsPrior p) -> HyperRep p
HyperRep forall a b. (a -> b) -> a -> b
$ forall p l. Conjugate p l => l -> Params p -> Support l -> Params p
updatePrior @p @l l
lk Hyper (AsPrior p)
pr Support l
val)
          r HyperRep
priors
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', r HyperRep
priors')
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support l
val
  sampleConst :: forall d.
(Distribution d, SampleCtx (UpdatePriorsI r) d) =>
String -> d -> Params d -> UpdatePriorsI r (Support d)
sampleConst String
_ d
_ Params d
_ = forall (r :: (* -> *) -> *) a.
StateT (Trace r, r HyperRep) Maybe a -> UpdatePriorsI r a
UpdatePriorsI forall a b. (a -> b) -> a -> b
$ do
    (Trace r
trace, r HyperRep
priors) <- forall s (m :: * -> *). MonadState s m => m s
get
    (Support d
val  , Trace r
trace') <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace
    forall s (m :: * -> *). MonadState s m => s -> m ()
put (Trace r
trace', r HyperRep
priors)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Support d
val
  permutationPlate :: forall a. Ord a => Int -> UpdatePriorsI r a -> UpdatePriorsI r [a]
permutationPlate = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM

getPosterior
  :: r HyperRep -> Trace r -> UpdatePriorsI r a -> Maybe (r HyperRep)
getPosterior :: forall (r :: (* -> *) -> *) a.
r HyperRep -> Trace r -> UpdatePriorsI r a -> Maybe (r HyperRep)
getPosterior r HyperRep
priors Trace r
trace (UpdatePriorsI StateT (Trace r, r HyperRep) Maybe a
model) = do
  (Trace r
_trace, r HyperRep
posteriors) <- forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT StateT (Trace r, r HyperRep) Maybe a
model (Trace r
trace, r HyperRep
priors)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure r HyperRep
posteriors

-- show how a trace is generated by a model
-- ----------------------------------------

newtype ShowTraceI r a = ShowTraceI (MaybeT (WriterT String (State (Trace r))) a)
  deriving (forall a b. a -> ShowTraceI r b -> ShowTraceI r a
forall a b. (a -> b) -> ShowTraceI r a -> ShowTraceI r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (r :: (* -> *) -> *) a b.
a -> ShowTraceI r b -> ShowTraceI r a
forall (r :: (* -> *) -> *) a b.
(a -> b) -> ShowTraceI r a -> ShowTraceI r b
<$ :: forall a b. a -> ShowTraceI r b -> ShowTraceI r a
$c<$ :: forall (r :: (* -> *) -> *) a b.
a -> ShowTraceI r b -> ShowTraceI r a
fmap :: forall a b. (a -> b) -> ShowTraceI r a -> ShowTraceI r b
$cfmap :: forall (r :: (* -> *) -> *) a b.
(a -> b) -> ShowTraceI r a -> ShowTraceI r b
Functor, forall a. a -> ShowTraceI r a
forall a b. ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r a
forall a b. ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
forall a b.
ShowTraceI r (a -> b) -> ShowTraceI r a -> ShowTraceI r b
forall a b c.
(a -> b -> c) -> ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (r :: (* -> *) -> *). Functor (ShowTraceI r)
forall (r :: (* -> *) -> *) a. a -> ShowTraceI r a
forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r a
forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
forall (r :: (* -> *) -> *) a b.
ShowTraceI r (a -> b) -> ShowTraceI r a -> ShowTraceI r b
forall (r :: (* -> *) -> *) a b c.
(a -> b -> c) -> ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r c
<* :: forall a b. ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r a
$c<* :: forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r a
*> :: forall a b. ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
$c*> :: forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
liftA2 :: forall a b c.
(a -> b -> c) -> ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r c
$cliftA2 :: forall (r :: (* -> *) -> *) a b c.
(a -> b -> c) -> ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r c
<*> :: forall a b.
ShowTraceI r (a -> b) -> ShowTraceI r a -> ShowTraceI r b
$c<*> :: forall (r :: (* -> *) -> *) a b.
ShowTraceI r (a -> b) -> ShowTraceI r a -> ShowTraceI r b
pure :: forall a. a -> ShowTraceI r a
$cpure :: forall (r :: (* -> *) -> *) a. a -> ShowTraceI r a
Applicative, forall a. a -> ShowTraceI r a
forall a b. ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
forall a b.
ShowTraceI r a -> (a -> ShowTraceI r b) -> ShowTraceI r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (r :: (* -> *) -> *). Applicative (ShowTraceI r)
forall (r :: (* -> *) -> *) a. a -> ShowTraceI r a
forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> (a -> ShowTraceI r b) -> ShowTraceI r b
return :: forall a. a -> ShowTraceI r a
$creturn :: forall (r :: (* -> *) -> *) a. a -> ShowTraceI r a
>> :: forall a b. ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
$c>> :: forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> ShowTraceI r b -> ShowTraceI r b
>>= :: forall a b.
ShowTraceI r a -> (a -> ShowTraceI r b) -> ShowTraceI r b
$c>>= :: forall (r :: (* -> *) -> *) a b.
ShowTraceI r a -> (a -> ShowTraceI r b) -> ShowTraceI r b
Monad)

showTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> ShowTraceI r (Support l)
showTraceItem :: forall l (r :: (* -> *) -> *).
(Show (Support l), Typeable l, Typeable (Support l)) =>
String -> ShowTraceI r (Support l)
showTraceItem String
name = forall (r :: (* -> *) -> *) a.
MaybeT (WriterT String (State (Trace r))) a -> ShowTraceI r a
ShowTraceI forall a b. (a -> b) -> a -> b
$ do
  Trace r
trace         <- forall s (m :: * -> *). MonadState s m => m s
get
  (Support l
val, Trace r
trace') <- forall (m :: * -> *) a. m (Maybe a) -> MaybeT m a
MaybeT forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace
  forall s (m :: * -> *). MonadState s m => s -> m ()
put Trace r
trace'
  let distName :: String
distName = forall a. Show a => a -> String
show (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall {k} (t :: k). Proxy t
Proxy :: Proxy l))
  forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell
    forall a b. (a -> b) -> a -> b
$  String
"Sampled value "
    forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Support l
val
    forall a. Semigroup a => a -> a -> a
<> String
" from a "
    forall a. Semigroup a => a -> a -> a
<> String
distName
    forall a. Semigroup a => a -> a -> a
<> String
" at "
    forall a. Semigroup a => a -> a -> a
<> String
name
    forall a. Semigroup a => a -> a -> a
<> String
".\n"
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Support l
val

instance RandomInterpreter (ShowTraceI r) r where
  type SampleCtx (ShowTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> ShowTraceI r (Support l)
  sampleValue :: forall p l.
(Conjugate p l, Typeable (Support l), Typeable l,
 Show (Support l)) =>
String -> l -> Accessor r p -> ShowTraceI r (Support l)
sampleValue String
name l
_ Accessor r p
acc = forall l (r :: (* -> *) -> *).
(Show (Support l), Typeable l, Typeable (Support l)) =>
String -> ShowTraceI r (Support l)
showTraceItem @l String
name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (ShowTraceI r) d)
    => String
    -> d
    -> Params d
    -> ShowTraceI r (Support d)
  sampleConst :: forall d.
(Distribution d, SampleCtx (ShowTraceI r) d) =>
String -> d -> Params d -> ShowTraceI r (Support d)
sampleConst String
name d
_ Params d
_ = forall l (r :: (* -> *) -> *).
(Show (Support l), Typeable l, Typeable (Support l)) =>
String -> ShowTraceI r (Support l)
showTraceItem @d String
name
  permutationPlate :: forall a. Ord a => Int -> ShowTraceI r a -> ShowTraceI r [a]
permutationPlate = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM

showTrace :: Trace r -> ShowTraceI r a -> (Maybe a, String)
showTrace :: forall (r :: (* -> *) -> *) a.
Trace r -> ShowTraceI r a -> (Maybe a, String)
showTrace Trace r
trace (ShowTraceI MaybeT (WriterT String (State (Trace r))) a
model) =
  forall s a. State s a -> s -> a
evalState (forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (forall (m :: * -> *) a. MaybeT m a -> m (Maybe a)
runMaybeT MaybeT (WriterT String (State (Trace r))) a
model)) Trace r
trace

printTrace :: Trace r -> ShowTraceI r a -> IO ()
printTrace :: forall (r :: (* -> *) -> *) a. Trace r -> ShowTraceI r a -> IO ()
printTrace Trace r
trace ShowTraceI r a
model = do
  let (Maybe a
res, String
txt) = forall (r :: (* -> *) -> *) a.
Trace r -> ShowTraceI r a -> (Maybe a, String)
showTrace Trace r
trace ShowTraceI r a
model
  String -> IO ()
putStrLn String
txt
  case Maybe a
res of
    Maybe a
Nothing -> do
      String -> IO ()
putStrLn String
"Trace does not match the model (stops here)"
    Just a
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- tracing interpreter (for debugging)
-- -----------------------------------

newtype TraceTraceI r a = TraceTraceI (State (Trace r) a)
  deriving (forall a b. a -> TraceTraceI r b -> TraceTraceI r a
forall a b. (a -> b) -> TraceTraceI r a -> TraceTraceI r b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (r :: (* -> *) -> *) a b.
a -> TraceTraceI r b -> TraceTraceI r a
forall (r :: (* -> *) -> *) a b.
(a -> b) -> TraceTraceI r a -> TraceTraceI r b
<$ :: forall a b. a -> TraceTraceI r b -> TraceTraceI r a
$c<$ :: forall (r :: (* -> *) -> *) a b.
a -> TraceTraceI r b -> TraceTraceI r a
fmap :: forall a b. (a -> b) -> TraceTraceI r a -> TraceTraceI r b
$cfmap :: forall (r :: (* -> *) -> *) a b.
(a -> b) -> TraceTraceI r a -> TraceTraceI r b
Functor, forall a. a -> TraceTraceI r a
forall a b. TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r a
forall a b. TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
forall a b.
TraceTraceI r (a -> b) -> TraceTraceI r a -> TraceTraceI r b
forall a b c.
(a -> b -> c)
-> TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall (r :: (* -> *) -> *). Functor (TraceTraceI r)
forall (r :: (* -> *) -> *) a. a -> TraceTraceI r a
forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r a
forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
forall (r :: (* -> *) -> *) a b.
TraceTraceI r (a -> b) -> TraceTraceI r a -> TraceTraceI r b
forall (r :: (* -> *) -> *) a b c.
(a -> b -> c)
-> TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r c
<* :: forall a b. TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r a
$c<* :: forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r a
*> :: forall a b. TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
$c*> :: forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
liftA2 :: forall a b c.
(a -> b -> c)
-> TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r c
$cliftA2 :: forall (r :: (* -> *) -> *) a b c.
(a -> b -> c)
-> TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r c
<*> :: forall a b.
TraceTraceI r (a -> b) -> TraceTraceI r a -> TraceTraceI r b
$c<*> :: forall (r :: (* -> *) -> *) a b.
TraceTraceI r (a -> b) -> TraceTraceI r a -> TraceTraceI r b
pure :: forall a. a -> TraceTraceI r a
$cpure :: forall (r :: (* -> *) -> *) a. a -> TraceTraceI r a
Applicative, forall a. a -> TraceTraceI r a
forall a b. TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
forall a b.
TraceTraceI r a -> (a -> TraceTraceI r b) -> TraceTraceI r b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
forall (r :: (* -> *) -> *). Applicative (TraceTraceI r)
forall (r :: (* -> *) -> *) a. a -> TraceTraceI r a
forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> (a -> TraceTraceI r b) -> TraceTraceI r b
return :: forall a. a -> TraceTraceI r a
$creturn :: forall (r :: (* -> *) -> *) a. a -> TraceTraceI r a
>> :: forall a b. TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
$c>> :: forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> TraceTraceI r b -> TraceTraceI r b
>>= :: forall a b.
TraceTraceI r a -> (a -> TraceTraceI r b) -> TraceTraceI r b
$c>>= :: forall (r :: (* -> *) -> *) a b.
TraceTraceI r a -> (a -> TraceTraceI r b) -> TraceTraceI r b
Monad)

traceTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> TraceTraceI r (Support l)
traceTraceItem :: forall l (r :: (* -> *) -> *).
(Show (Support l), Typeable l, Typeable (Support l)) =>
String -> TraceTraceI r (Support l)
traceTraceItem String
name = forall (r :: (* -> *) -> *) a. State (Trace r) a -> TraceTraceI r a
TraceTraceI forall a b. (a -> b) -> a -> b
$ do
  Trace r
trace <- forall s (m :: * -> *). MonadState s m => m s
get
  let loc :: String
loc = forall a. Show a => a -> String
show (forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (forall {k} (t :: k). Proxy t
Proxy :: Proxy l)) forall a. Semigroup a => a -> a -> a
<> String
" at " forall a. Semigroup a => a -> a -> a
<> String
name
  case forall a (r :: (* -> *) -> *).
Typeable a =>
Trace r -> Maybe (a, Trace r)
takeTrace Trace r
trace of
    Just (Support l
val, Trace r
trace') -> do
      forall s (m :: * -> *). MonadState s m => s -> m ()
put Trace r
trace'
      forall (f :: * -> *). Applicative f => String -> f ()
DT.traceM forall a b. (a -> b) -> a -> b
$ String
"Sampled value " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Support l
val forall a. Semigroup a => a -> a -> a
<> String
" from a " forall a. Semigroup a => a -> a -> a
<> String
loc forall a. Semigroup a => a -> a -> a
<> String
"."
      forall (f :: * -> *) a. Applicative f => a -> f a
pure Support l
val
    Maybe (Support l, Trace r)
Nothing -> forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Incompatible trace at " forall a. Semigroup a => a -> a -> a
<> String
loc forall a. Semigroup a => a -> a -> a
<> String
"."

instance RandomInterpreter (TraceTraceI r) r where
  type SampleCtx (TraceTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceTraceI r (Support l)
  sampleValue :: forall p l.
(Conjugate p l, Typeable (Support l), Typeable l,
 Show (Support l)) =>
String -> l -> Accessor r p -> TraceTraceI r (Support l)
sampleValue String
name l
_ Accessor r p
acc = forall l (r :: (* -> *) -> *).
(Show (Support l), Typeable l, Typeable (Support l)) =>
String -> TraceTraceI r (Support l)
traceTraceItem @l String
name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (TraceTraceI r) d)
    => String
    -> d
    -> Params d
    -> TraceTraceI r (Support d)
  sampleConst :: forall d.
(Distribution d, SampleCtx (TraceTraceI r) d) =>
String -> d -> Params d -> TraceTraceI r (Support d)
sampleConst String
name d
_ Params d
_ = forall l (r :: (* -> *) -> *).
(Show (Support l), Typeable l, Typeable (Support l)) =>
String -> TraceTraceI r (Support l)
traceTraceItem @d String
name
  permutationPlate :: forall a. Ord a => Int -> TraceTraceI r a -> TraceTraceI r [a]
permutationPlate = forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM

traceTrace :: Trace r -> TraceTraceI r a -> a
traceTrace :: forall (r :: (* -> *) -> *) a. Trace r -> TraceTraceI r a -> a
traceTrace Trace r
trace (TraceTraceI State (Trace r) a
model) = forall s a. State s a -> s -> a
evalState State (Trace r) a
model Trace r
trace

-------------------
-- distributions --
-------------------

-- Beta
-- ----

data Beta = Beta
  deriving (Beta -> Beta -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Beta -> Beta -> Bool
$c/= :: Beta -> Beta -> Bool
== :: Beta -> Beta -> Bool
$c== :: Beta -> Beta -> Bool
Eq, Eq Beta
Beta -> Beta -> Bool
Beta -> Beta -> Ordering
Beta -> Beta -> Beta
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Beta -> Beta -> Beta
$cmin :: Beta -> Beta -> Beta
max :: Beta -> Beta -> Beta
$cmax :: Beta -> Beta -> Beta
>= :: Beta -> Beta -> Bool
$c>= :: Beta -> Beta -> Bool
> :: Beta -> Beta -> Bool
$c> :: Beta -> Beta -> Bool
<= :: Beta -> Beta -> Bool
$c<= :: Beta -> Beta -> Bool
< :: Beta -> Beta -> Bool
$c< :: Beta -> Beta -> Bool
compare :: Beta -> Beta -> Ordering
$ccompare :: Beta -> Beta -> Ordering
Ord, Int -> Beta -> ShowS
[Beta] -> ShowS
Beta -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Beta] -> ShowS
$cshowList :: [Beta] -> ShowS
show :: Beta -> String
$cshow :: Beta -> String
showsPrec :: Int -> Beta -> ShowS
$cshowsPrec :: Int -> Beta -> ShowS
Show, forall x. Rep Beta x -> Beta
forall x. Beta -> Rep Beta x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Beta x -> Beta
$cfrom :: forall x. Beta -> Rep Beta x
Generic)

instance Distribution Beta where
  type Params Beta = (Double, Double)
  type Support Beta = Double
  distSample :: forall (m :: * -> *).
PrimMonad m =>
Beta -> Params Beta -> Prob m (Support Beta)
distSample Beta
_ = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> Prob m Double
beta
  distLogP :: Beta -> Params Beta -> Support Beta -> Double
distLogP Beta
_ (Double
a, Double
b) Support Beta
p =
    forall a. Floating a => a -> a
log (Support Beta
p forall a. Floating a => a -> a -> a
** (Double
a forall a. Num a => a -> a -> a
- Double
1)) forall a. Num a => a -> a -> a
+ forall a. Floating a => a -> a
log ((Double
1 forall a. Num a => a -> a -> a
- Support Beta
p) forall a. Floating a => a -> a -> a
** (Double
b forall a. Num a => a -> a -> a
- Double
1)) forall a. Num a => a -> a -> a
- Double -> Double -> Double
logBeta Double
a Double
b

instance Jeffreys (AsPrior Beta) where
  jeffreysPrior :: Hyper (AsPrior Beta)
jeffreysPrior = (Double
0.5, Double
0.5)

instance Uniform (AsPrior Beta) where
  uniformPrior :: Hyper (AsPrior Beta)
uniformPrior = (Double
1, Double
1)

instance Prior (AsPrior Beta) where
  sampleProbs :: forall (m :: * -> *).
PrimMonad m =>
Hyper (AsPrior Beta) -> Prob m (Probs (AsPrior Beta))
sampleProbs = forall a (m :: * -> *).
(Distribution a, PrimMonad m) =>
a -> Params a -> Prob m (Support a)
distSample Beta
Beta

-- Bernoulli
-- ---------

data Bernoulli = Bernoulli
  deriving (Bernoulli -> Bernoulli -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Bernoulli -> Bernoulli -> Bool
$c/= :: Bernoulli -> Bernoulli -> Bool
== :: Bernoulli -> Bernoulli -> Bool
$c== :: Bernoulli -> Bernoulli -> Bool
Eq, Eq Bernoulli
Bernoulli -> Bernoulli -> Bool
Bernoulli -> Bernoulli -> Ordering
Bernoulli -> Bernoulli -> Bernoulli
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Bernoulli -> Bernoulli -> Bernoulli
$cmin :: Bernoulli -> Bernoulli -> Bernoulli
max :: Bernoulli -> Bernoulli -> Bernoulli
$cmax :: Bernoulli -> Bernoulli -> Bernoulli
>= :: Bernoulli -> Bernoulli -> Bool
$c>= :: Bernoulli -> Bernoulli -> Bool
> :: Bernoulli -> Bernoulli -> Bool
$c> :: Bernoulli -> Bernoulli -> Bool
<= :: Bernoulli -> Bernoulli -> Bool
$c<= :: Bernoulli -> Bernoulli -> Bool
< :: Bernoulli -> Bernoulli -> Bool
$c< :: Bernoulli -> Bernoulli -> Bool
compare :: Bernoulli -> Bernoulli -> Ordering
$ccompare :: Bernoulli -> Bernoulli -> Ordering
Ord, Int -> Bernoulli -> ShowS
[Bernoulli] -> ShowS
Bernoulli -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Bernoulli] -> ShowS
$cshowList :: [Bernoulli] -> ShowS
show :: Bernoulli -> String
$cshow :: Bernoulli -> String
showsPrec :: Int -> Bernoulli -> ShowS
$cshowsPrec :: Int -> Bernoulli -> ShowS
Show, forall x. Rep Bernoulli x -> Bernoulli
forall x. Bernoulli -> Rep Bernoulli x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Bernoulli x -> Bernoulli
$cfrom :: forall x. Bernoulli -> Rep Bernoulli x
Generic)

instance Distribution Bernoulli where
  type Params Bernoulli = Double
  type Support Bernoulli = Bool
  distSample :: forall (m :: * -> *).
PrimMonad m =>
Bernoulli -> Params Bernoulli -> Prob m (Support Bernoulli)
distSample Bernoulli
_ = forall (m :: * -> *). PrimMonad m => Double -> Prob m Bool
bernoulli
  distLogP :: Bernoulli -> Params Bernoulli -> Support Bernoulli -> Double
distLogP Bernoulli
_ Params Bernoulli
p Bool
Support Bernoulli
True  = forall a. Floating a => a -> a
log Params Bernoulli
p
  distLogP Bernoulli
_ Params Bernoulli
p Bool
Support Bernoulli
False = forall a. Floating a => a -> a
log (Double
1 forall a. Num a => a -> a -> a
- Params Bernoulli
p)

-- Binomial
-- --------

newtype Binomial = Binomial Int
  deriving (Binomial -> Binomial -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Binomial -> Binomial -> Bool
$c/= :: Binomial -> Binomial -> Bool
== :: Binomial -> Binomial -> Bool
$c== :: Binomial -> Binomial -> Bool
Eq, Eq Binomial
Binomial -> Binomial -> Bool
Binomial -> Binomial -> Ordering
Binomial -> Binomial -> Binomial
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Binomial -> Binomial -> Binomial
$cmin :: Binomial -> Binomial -> Binomial
max :: Binomial -> Binomial -> Binomial
$cmax :: Binomial -> Binomial -> Binomial
>= :: Binomial -> Binomial -> Bool
$c>= :: Binomial -> Binomial -> Bool
> :: Binomial -> Binomial -> Bool
$c> :: Binomial -> Binomial -> Bool
<= :: Binomial -> Binomial -> Bool
$c<= :: Binomial -> Binomial -> Bool
< :: Binomial -> Binomial -> Bool
$c< :: Binomial -> Binomial -> Bool
compare :: Binomial -> Binomial -> Ordering
$ccompare :: Binomial -> Binomial -> Ordering
Ord, Int -> Binomial -> ShowS
[Binomial] -> ShowS
Binomial -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Binomial] -> ShowS
$cshowList :: [Binomial] -> ShowS
show :: Binomial -> String
$cshow :: Binomial -> String
showsPrec :: Int -> Binomial -> ShowS
$cshowsPrec :: Int -> Binomial -> ShowS
Show, forall x. Rep Binomial x -> Binomial
forall x. Binomial -> Rep Binomial x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Binomial x -> Binomial
$cfrom :: forall x. Binomial -> Rep Binomial x
Generic)

instance Distribution Binomial where
  type Params Binomial = Double
  type Support Binomial = Int
  distSample :: forall (m :: * -> *).
PrimMonad m =>
Binomial -> Params Binomial -> Prob m (Support Binomial)
distSample (Binomial Int
n) = forall (m :: * -> *). PrimMonad m => Int -> Double -> Prob m Int
binomial Int
n
  distLogP :: Binomial -> Params Binomial -> Support Binomial -> Double
distLogP (Binomial Int
n) Params Binomial
p Support Binomial
k =
    Int -> Int -> Double
logChoose Int
n Support Binomial
k forall a. Num a => a -> a -> a
+ (forall a. Floating a => a -> a
log Params Binomial
p forall a. Num a => a -> a -> a
* Double
k') forall a. Num a => a -> a -> a
+ (forall a. Floating a => a -> a
log (Double
1 forall a. Num a => a -> a -> a
- Params Binomial
p) forall a. Num a => a -> a -> a
* (Double
n' forall a. Num a => a -> a -> a
- Double
k'))
   where
    k' :: Double
k' = Int -> Double
int2Double Support Binomial
k
    n' :: Double
n' = Int -> Double
int2Double Int
n

-- Categorical
-- -----------

type Categorical :: Nat -> Type
data Categorical n = Categorical
  deriving (Categorical n -> Categorical n -> Bool
forall (n :: Nat). Categorical n -> Categorical n -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Categorical n -> Categorical n -> Bool
$c/= :: forall (n :: Nat). Categorical n -> Categorical n -> Bool
== :: Categorical n -> Categorical n -> Bool
$c== :: forall (n :: Nat). Categorical n -> Categorical n -> Bool
Eq, Categorical n -> Categorical n -> Bool
Categorical n -> Categorical n -> Ordering
forall (n :: Nat). Eq (Categorical n)
forall (n :: Nat). Categorical n -> Categorical n -> Bool
forall (n :: Nat). Categorical n -> Categorical n -> Ordering
forall (n :: Nat). Categorical n -> Categorical n -> Categorical n
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Categorical n -> Categorical n -> Categorical n
$cmin :: forall (n :: Nat). Categorical n -> Categorical n -> Categorical n
max :: Categorical n -> Categorical n -> Categorical n
$cmax :: forall (n :: Nat). Categorical n -> Categorical n -> Categorical n
>= :: Categorical n -> Categorical n -> Bool
$c>= :: forall (n :: Nat). Categorical n -> Categorical n -> Bool
> :: Categorical n -> Categorical n -> Bool
$c> :: forall (n :: Nat). Categorical n -> Categorical n -> Bool
<= :: Categorical n -> Categorical n -> Bool
$c<= :: forall (n :: Nat). Categorical n -> Categorical n -> Bool
< :: Categorical n -> Categorical n -> Bool
$c< :: forall (n :: Nat). Categorical n -> Categorical n -> Bool
compare :: Categorical n -> Categorical n -> Ordering
$ccompare :: forall (n :: Nat). Categorical n -> Categorical n -> Ordering
Ord, Int -> Categorical n -> ShowS
forall (n :: Nat). Int -> Categorical n -> ShowS
forall (n :: Nat). [Categorical n] -> ShowS
forall (n :: Nat). Categorical n -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Categorical n] -> ShowS
$cshowList :: forall (n :: Nat). [Categorical n] -> ShowS
show :: Categorical n -> String
$cshow :: forall (n :: Nat). Categorical n -> String
showsPrec :: Int -> Categorical n -> ShowS
$cshowsPrec :: forall (n :: Nat). Int -> Categorical n -> ShowS
Show, forall (n :: Nat) x. Rep (Categorical n) x -> Categorical n
forall (n :: Nat) x. Categorical n -> Rep (Categorical n) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall (n :: Nat) x. Rep (Categorical n) x -> Categorical n
$cfrom :: forall (n :: Nat) x. Categorical n -> Rep (Categorical n) x
Generic)

instance Distribution (Categorical n) where
  type Params (Categorical n) = V.Vector Double
  type Support (Categorical n) = Int
  distSample :: forall (m :: * -> *).
PrimMonad m =>
Categorical n
-> Params (Categorical n) -> Prob m (Support (Categorical n))
distSample Categorical n
_ = forall (f :: * -> *) (m :: * -> *).
(Foldable f, PrimMonad m) =>
f Double -> Prob m Int
categorical
  distLogP :: Categorical n
-> Params (Categorical n) -> Support (Categorical n) -> Double
distLogP Categorical n
_ Params (Categorical n)
ps Support (Categorical n)
cat = forall a. Floating a => a -> a
log forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe Double
0 forall a b. (a -> b) -> a -> b
$ Params (Categorical n)
ps forall a. Vector a -> Int -> Maybe a
V.!? Support (Categorical n)
cat

-- Dirichlet
-- ---------

type Dirichlet :: Nat -> Type
data Dirichlet n = Dirichlet
  deriving (Dirichlet n -> Dirichlet n -> Bool
forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Dirichlet n -> Dirichlet n -> Bool
$c/= :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
== :: Dirichlet n -> Dirichlet n -> Bool
$c== :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
Eq, Dirichlet n -> Dirichlet n -> Bool
Dirichlet n -> Dirichlet n -> Ordering
forall (n :: Nat). Eq (Dirichlet n)
forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
forall (n :: Nat). Dirichlet n -> Dirichlet n -> Ordering
forall (n :: Nat). Dirichlet n -> Dirichlet n -> Dirichlet n
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Dirichlet n -> Dirichlet n -> Dirichlet n
$cmin :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Dirichlet n
max :: Dirichlet n -> Dirichlet n -> Dirichlet n
$cmax :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Dirichlet n
>= :: Dirichlet n -> Dirichlet n -> Bool
$c>= :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
> :: Dirichlet n -> Dirichlet n -> Bool
$c> :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
<= :: Dirichlet n -> Dirichlet n -> Bool
$c<= :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
< :: Dirichlet n -> Dirichlet n -> Bool
$c< :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Bool
compare :: Dirichlet n -> Dirichlet n -> Ordering
$ccompare :: forall (n :: Nat). Dirichlet n -> Dirichlet n -> Ordering
Ord, Int -> Dirichlet n -> ShowS
forall (n :: Nat). Int -> Dirichlet n -> ShowS
forall (n :: Nat). [Dirichlet n] -> ShowS
forall (n :: Nat). Dirichlet n -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Dirichlet n] -> ShowS
$cshowList :: forall (n :: Nat). [Dirichlet n] -> ShowS
show :: Dirichlet n -> String
$cshow :: forall (n :: Nat). Dirichlet n -> String
showsPrec :: Int -> Dirichlet n -> ShowS
$cshowsPrec :: forall (n :: Nat). Int -> Dirichlet n -> ShowS
Show, forall (n :: Nat) x. Rep (Dirichlet n) x -> Dirichlet n
forall (n :: Nat) x. Dirichlet n -> Rep (Dirichlet n) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall (n :: Nat) x. Rep (Dirichlet n) x -> Dirichlet n
$cfrom :: forall (n :: Nat) x. Dirichlet n -> Rep (Dirichlet n) x
Generic)

instance Distribution (Dirichlet n) where
  type Params (Dirichlet n) = V.Vector Double
  type Support (Dirichlet n) = V.Vector Double
  distSample :: forall (m :: * -> *).
PrimMonad m =>
Dirichlet n
-> Params (Dirichlet n) -> Prob m (Support (Dirichlet n))
distSample Dirichlet n
_ = forall (f :: * -> *) (m :: * -> *).
(Traversable f, PrimMonad m) =>
f Double -> Prob m (f Double)
dirichlet
  distLogP :: Dirichlet n
-> Params (Dirichlet n) -> Support (Dirichlet n) -> Double
distLogP Dirichlet n
_ Params (Dirichlet n)
counts Support (Dirichlet n)
probs = Double
logp forall a. Num a => a -> a -> a
+ Double
logz
   where
    logp :: Double
logp = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith (\Double
a Double
x -> forall a. Floating a => a -> a
log Double
x forall a. Num a => a -> a -> a
* (Double
a forall a. Num a => a -> a -> a
- Double
1)) Params (Dirichlet n)
counts Support (Dirichlet n)
probs)
    logz :: Double
logz = Double -> Double
logGamma (forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Params (Dirichlet n)
counts) forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Double -> Double
logGamma forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Params (Dirichlet n)
counts)

instance KnownNat n => Jeffreys (AsPrior (Dirichlet n)) where
  jeffreysPrior :: Hyper (AsPrior (Dirichlet n))
jeffreysPrior = forall a. Int -> a -> Vector a
V.replicate (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal (forall {k} (t :: k). Proxy t
Proxy :: Proxy n)) Double
0.5

instance KnownNat n => Uniform (AsPrior (Dirichlet n)) where
  uniformPrior :: Hyper (AsPrior (Dirichlet n))
uniformPrior = forall a. Int -> a -> Vector a
V.replicate (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal (forall {k} (t :: k). Proxy t
Proxy :: Proxy n)) Double
1

instance Prior (AsPrior (Dirichlet n)) where
  sampleProbs :: forall (m :: * -> *).
PrimMonad m =>
Hyper (AsPrior (Dirichlet n))
-> Prob m (Probs (AsPrior (Dirichlet n)))
sampleProbs = forall a (m :: * -> *).
(Distribution a, PrimMonad m) =>
a -> Params a -> Prob m (Support a)
distSample forall (n :: Nat). Dirichlet n
Dirichlet

-- Geometric (from 0)
-- ------------------

data Geometric0 = Geometric0
  deriving (Geometric0 -> Geometric0 -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Geometric0 -> Geometric0 -> Bool
$c/= :: Geometric0 -> Geometric0 -> Bool
== :: Geometric0 -> Geometric0 -> Bool
$c== :: Geometric0 -> Geometric0 -> Bool
Eq, Eq Geometric0
Geometric0 -> Geometric0 -> Bool
Geometric0 -> Geometric0 -> Ordering
Geometric0 -> Geometric0 -> Geometric0
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Geometric0 -> Geometric0 -> Geometric0
$cmin :: Geometric0 -> Geometric0 -> Geometric0
max :: Geometric0 -> Geometric0 -> Geometric0
$cmax :: Geometric0 -> Geometric0 -> Geometric0
>= :: Geometric0 -> Geometric0 -> Bool
$c>= :: Geometric0 -> Geometric0 -> Bool
> :: Geometric0 -> Geometric0 -> Bool
$c> :: Geometric0 -> Geometric0 -> Bool
<= :: Geometric0 -> Geometric0 -> Bool
$c<= :: Geometric0 -> Geometric0 -> Bool
< :: Geometric0 -> Geometric0 -> Bool
$c< :: Geometric0 -> Geometric0 -> Bool
compare :: Geometric0 -> Geometric0 -> Ordering
$ccompare :: Geometric0 -> Geometric0 -> Ordering
Ord, Int -> Geometric0 -> ShowS
[Geometric0] -> ShowS
Geometric0 -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Geometric0] -> ShowS
$cshowList :: [Geometric0] -> ShowS
show :: Geometric0 -> String
$cshow :: Geometric0 -> String
showsPrec :: Int -> Geometric0 -> ShowS
$cshowsPrec :: Int -> Geometric0 -> ShowS
Show, forall x. Rep Geometric0 x -> Geometric0
forall x. Geometric0 -> Rep Geometric0 x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Geometric0 x -> Geometric0
$cfrom :: forall x. Geometric0 -> Rep Geometric0 x
Generic)

instance Distribution Geometric0 where
  type Params Geometric0 = Double
  type Support Geometric0 = Int
  distSample :: forall (m :: * -> *).
PrimMonad m =>
Geometric0 -> Params Geometric0 -> Prob m (Support Geometric0)
distSample Geometric0
_ = forall {m :: * -> *} {b}.
(PrimMonad m, Num b) =>
Double -> Prob m b
geometric0
   where
    geometric0 :: Double -> Prob m b
geometric0 Double
p = do
      Bool
coin <- forall (m :: * -> *). PrimMonad m => Double -> Prob m Bool
bernoulli Double
p
      if Bool
coin then forall (f :: * -> *) a. Applicative f => a -> f a
pure b
0 else (b
1 forall a. Num a => a -> a -> a
+) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Prob m b
geometric0 Double
p
  distLogP :: Geometric0 -> Params Geometric0 -> Support Geometric0 -> Double
distLogP Geometric0
_ Params Geometric0
p Support Geometric0
val | Support Geometric0
val forall a. Ord a => a -> a -> Bool
>= Int
0  = (forall a. Floating a => a -> a
log (Double
1 forall a. Num a => a -> a -> a
- Params Geometric0
p) forall a. Num a => a -> a -> a
* Int -> Double
int2Double Support Geometric0
val) forall a. Num a => a -> a -> a
+ forall a. Floating a => a -> a
log Params Geometric0
p
                   | Bool
otherwise = forall a. Floating a => a -> a
log Double
0

-- Geometric (from 1)
-- ------------------

data Geometric1 = Geometric1
  deriving (Geometric1 -> Geometric1 -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Geometric1 -> Geometric1 -> Bool
$c/= :: Geometric1 -> Geometric1 -> Bool
== :: Geometric1 -> Geometric1 -> Bool
$c== :: Geometric1 -> Geometric1 -> Bool
Eq, Eq Geometric1
Geometric1 -> Geometric1 -> Bool
Geometric1 -> Geometric1 -> Ordering
Geometric1 -> Geometric1 -> Geometric1
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Geometric1 -> Geometric1 -> Geometric1
$cmin :: Geometric1 -> Geometric1 -> Geometric1
max :: Geometric1 -> Geometric1 -> Geometric1
$cmax :: Geometric1 -> Geometric1 -> Geometric1
>= :: Geometric1 -> Geometric1 -> Bool
$c>= :: Geometric1 -> Geometric1 -> Bool
> :: Geometric1 -> Geometric1 -> Bool
$c> :: Geometric1 -> Geometric1 -> Bool
<= :: Geometric1 -> Geometric1 -> Bool
$c<= :: Geometric1 -> Geometric1 -> Bool
< :: Geometric1 -> Geometric1 -> Bool
$c< :: Geometric1 -> Geometric1 -> Bool
compare :: Geometric1 -> Geometric1 -> Ordering
$ccompare :: Geometric1 -> Geometric1 -> Ordering
Ord, Int -> Geometric1 -> ShowS
[Geometric1] -> ShowS
Geometric1 -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Geometric1] -> ShowS
$cshowList :: [Geometric1] -> ShowS
show :: Geometric1 -> String
$cshow :: Geometric1 -> String
showsPrec :: Int -> Geometric1 -> ShowS
$cshowsPrec :: Int -> Geometric1 -> ShowS
Show, forall x. Rep Geometric1 x -> Geometric1
forall x. Geometric1 -> Rep Geometric1 x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Geometric1 x -> Geometric1
$cfrom :: forall x. Geometric1 -> Rep Geometric1 x
Generic)

instance Distribution Geometric1 where
  type Params Geometric1 = Double
  type Support Geometric1 = Int
  distSample :: forall (m :: * -> *).
PrimMonad m =>
Geometric1 -> Params Geometric1 -> Prob m (Support Geometric1)
distSample Geometric1
_ = forall {m :: * -> *} {b}.
(PrimMonad m, Num b) =>
Double -> Prob m b
geometric1
   where
    geometric1 :: Double -> Prob m b
geometric1 Double
p = do
      Bool
coin <- forall (m :: * -> *). PrimMonad m => Double -> Prob m Bool
bernoulli Double
p
      if Bool
coin then forall (f :: * -> *) a. Applicative f => a -> f a
pure b
1 else (b
1 forall a. Num a => a -> a -> a
+) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Prob m b
geometric1 Double
p
  distLogP :: Geometric1 -> Params Geometric1 -> Support Geometric1 -> Double
distLogP Geometric1
_ Params Geometric1
p Support Geometric1
val | Support Geometric1
val forall a. Ord a => a -> a -> Bool
>= Int
1  = (forall a. Floating a => a -> a
log (Double
1 forall a. Num a => a -> a -> a
- Params Geometric1
p) forall a. Num a => a -> a -> a
* Int -> Double
int2Double (Support Geometric1
val forall a. Num a => a -> a -> a
- Int
1)) forall a. Num a => a -> a -> a
+ forall a. Floating a => a -> a
log Params Geometric1
p
                   | Bool
otherwise = forall a. Floating a => a -> a
log Double
0

-----------------------------
-- conjugate distributions --
-----------------------------

-- beta bernoulli
-- --------------

instance Conjugate Beta Bernoulli where
  priorSingleton :: Beta
priorSingleton = Beta
Beta
  updatePrior :: Bernoulli -> Params Beta -> Support Bernoulli -> Params Beta
updatePrior Bernoulli
_ (Double
a, Double
b) Bool
Support Bernoulli
False = (Double
a, Double
b forall a. Num a => a -> a -> a
+ Double
1)
  updatePrior Bernoulli
_ (Double
a, Double
b) Bool
Support Bernoulli
True  = (Double
a forall a. Num a => a -> a -> a
+ Double
1, Double
b)
  predLogP :: Bernoulli -> Params Beta -> Support Bernoulli -> Double
predLogP Bernoulli
_ (Double
a, Double
b) Bool
Support Bernoulli
False = forall a. Floating a => a -> a
log Double
a forall a. Num a => a -> a -> a
- forall a. Floating a => a -> a
log (Double
a forall a. Num a => a -> a -> a
+ Double
b)
  predLogP Bernoulli
_ (Double
a, Double
b) Bool
Support Bernoulli
True  = forall a. Floating a => a -> a
log Double
b forall a. Num a => a -> a -> a
- forall a. Floating a => a -> a
log (Double
a forall a. Num a => a -> a -> a
+ Double
b)

-- B(1+a, b)   G(1+a) G(b) G(a+b)   G(1+a) G(a+b)   a G(a) G(a+b)        a
-- --------- = ------------------ = ------------- = ----------------- = ---
-- B(a, b)     G(a) G(b) G(a+b+1)   G(a) G(a+b+1)   G(a) (a+b) G(a+b)   a+b

-- beta binomial
-- -------------

instance Conjugate Beta Binomial where
  priorSingleton :: Beta
priorSingleton = Beta
Beta
  updatePrior :: Binomial -> Params Beta -> Support Binomial -> Params Beta
updatePrior (Binomial Int
n) (Double
a, Double
b) Support Binomial
x = (Double
a forall a. Num a => a -> a -> a
+ Double
x', Double
b forall a. Num a => a -> a -> a
+ (Double
n' forall a. Num a => a -> a -> a
- Double
x'))
   where
    x' :: Double
x' = Int -> Double
int2Double Support Binomial
x
    n' :: Double
n' = Int -> Double
int2Double Int
n
  predLogP :: Binomial -> Params Beta -> Support Binomial -> Double
predLogP (Binomial Int
n) (Double
a, Double
b) Support Binomial
k =
    Int -> Int -> Double
logChoose Int
n Support Binomial
k forall a. Num a => a -> a -> a
+ Double -> Double -> Double
logBeta (Double
k' forall a. Num a => a -> a -> a
+ Double
a) (Double
n' forall a. Num a => a -> a -> a
- Double
k' forall a. Num a => a -> a -> a
+ Double
a) forall a. Num a => a -> a -> a
- Double -> Double -> Double
logBeta Double
a Double
b
   where
    n' :: Double
n' = Int -> Double
int2Double Int
n
    k' :: Double
k' = Int -> Double
int2Double Support Binomial
k

-- beta geometric0
-- ---------------

instance Conjugate Beta Geometric0 where
  priorSingleton :: Beta
priorSingleton = Beta
Beta
  updatePrior :: Geometric0 -> Params Beta -> Support Geometric0 -> Params Beta
updatePrior Geometric0
_ (Double
a, Double
b) Support Geometric0
k = (Double
a forall a. Num a => a -> a -> a
+ Double
1, Double
b forall a. Num a => a -> a -> a
+ Int -> Double
int2Double Support Geometric0
k)
  predLogP :: Geometric0 -> Params Beta -> Support Geometric0 -> Double
predLogP Geometric0
_ (Double
a, Double
b) Support Geometric0
k =
    (forall a. Floating a => a -> a
log Double
a forall a. Num a => a -> a -> a
+ Double -> Double
logGamma (Double
a forall a. Num a => a -> a -> a
+ Double
b) forall a. Num a => a -> a -> a
+ Double -> Double
logGamma (Double
k' forall a. Num a => a -> a -> a
+ Double
b))
      forall a. Num a => a -> a -> a
- (Double -> Double
logGamma Double
b forall a. Num a => a -> a -> a
+ Double -> Double
logGamma (Double
a forall a. Num a => a -> a -> a
+ Double
b forall a. Num a => a -> a -> a
+ Double
k' forall a. Num a => a -> a -> a
+ Double
1))
    where k' :: Double
k' = Int -> Double
int2Double Support Geometric0
k


-- beta geometric1
-- ---------------

instance Conjugate Beta Geometric1 where
  priorSingleton :: Beta
priorSingleton = Beta
Beta
  updatePrior :: Geometric1 -> Params Beta -> Support Geometric1 -> Params Beta
updatePrior Geometric1
_ (Double
a, Double
b) Support Geometric1
k = (Double
a forall a. Num a => a -> a -> a
+ Double
1, Double
b forall a. Num a => a -> a -> a
+ Int -> Double
int2Double (Support Geometric1
k forall a. Num a => a -> a -> a
- Int
1))
  predLogP :: Geometric1 -> Params Beta -> Support Geometric1 -> Double
predLogP Geometric1
_ (Double
a, Double
b) Support Geometric1
k =
    (forall a. Floating a => a -> a
log Double
a forall a. Num a => a -> a -> a
+ Double -> Double
logGamma (Double
a forall a. Num a => a -> a -> a
+ Double
b) forall a. Num a => a -> a -> a
+ Double -> Double
logGamma (Double
k' forall a. Num a => a -> a -> a
+ Double
b))
      forall a. Num a => a -> a -> a
- (Double -> Double
logGamma Double
b forall a. Num a => a -> a -> a
+ Double -> Double
logGamma (Double
a forall a. Num a => a -> a -> a
+ Double
b forall a. Num a => a -> a -> a
+ Double
k' forall a. Num a => a -> a -> a
+ Double
1))
    where k' :: Double
k' = Int -> Double
int2Double Support Geometric1
k forall a. Num a => a -> a -> a
- Double
1

-- dirichlet categorical
-- ---------------------

instance Conjugate (Dirichlet n) (Categorical n) where
  priorSingleton :: Dirichlet n
priorSingleton = forall (n :: Nat). Dirichlet n
Dirichlet
  updatePrior :: Categorical n
-> Params (Dirichlet n)
-> Support (Categorical n)
-> Params (Dirichlet n)
updatePrior Categorical n
_ Params (Dirichlet n)
counts Support (Categorical n)
obs
    | Support (Categorical n)
obs forall a. Ord a => a -> a -> Bool
>= Int
0, Support (Categorical n)
obs forall a. Ord a => a -> a -> Bool
< forall a. Vector a -> Int
V.length Params (Dirichlet n)
counts
    = Params (Dirichlet n)
counts forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Support (Categorical n)
obs, (Params (Dirichlet n)
counts forall a. Vector a -> Int -> a
V.! Support (Categorical n)
obs) forall a. Num a => a -> a -> a
+ Double
1)]
    | Bool
otherwise
    = Params (Dirichlet n)
counts
  predLogP :: Categorical n
-> Params (Dirichlet n) -> Support (Categorical n) -> Double
predLogP Categorical n
_ Params (Dirichlet n)
counts Support (Categorical n)
obs =
    forall a. Floating a => a -> a
log (forall a. a -> Maybe a -> a
fromMaybe Double
0 forall a b. (a -> b) -> a -> b
$ Params (Dirichlet n)
counts forall a. Vector a -> Int -> Maybe a
V.!? Support (Categorical n)
obs) forall a. Num a => a -> a -> a
- forall a. Floating a => a -> a
log (forall a. Num a => Vector a -> a
V.sum Params (Dirichlet n)
counts)