{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -Wno-deprecations #-}

-- |
-- Module      : Control.Monad.Bayes.Population
-- Description : Representation of distributions using multiple samples
-- Copyright   : (c) Adam Scibior, 2015-2020
-- License     : MIT
-- Maintainer  : leonhard.markert@tweag.io
-- Stability   : experimental
-- Portability : GHC
--
-- 'PopulationT' turns a single sample into a collection of weighted samples.
module Control.Monad.Bayes.Population
  ( PopulationT (..),
    runPopulationT,
    explicitPopulation,
    fromWeightedList,
    spawn,
    multinomial,
    resampleMultinomial,
    systematic,
    resampleSystematic,
    stratified,
    resampleStratified,
    extractEvidence,
    pushEvidence,
    proper,
    evidence,
    hoist,
    collapse,
    popAvg,
    withParticles,
  )
where

import Control.Applicative (Alternative)
import Control.Arrow (second)
import Control.Monad (forM, replicateM)
import Control.Monad.Bayes.Class
  ( MonadDistribution (..),
    MonadFactor (..),
    MonadMeasure,
    factor,
  )
import Control.Monad.Bayes.Weighted
  ( WeightedT,
    applyWeight,
    extractWeight,
    runWeightedT,
    weightedT,
  )
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import Data.Functor.Compose
import Data.List (unfoldr)
import Data.List qualified
import Data.Maybe (catMaybes)
import Data.Vector ((!))
import Data.Vector qualified as V
import Numeric.Log (Log, ln, sum)
import Numeric.Log qualified as Log
import Prelude hiding (all, sum)

-- | The old-fashioned, broken list transformer, adding a list/nondeterminism/choice effect.
--   It is not a valid monad transformer, but it is a valid 'Applicative'.
newtype ListT m a = ListT {forall (m :: * -> *) a. ListT m a -> Compose m [] a
getListT :: Compose m [] a}
  deriving newtype ((forall a b. (a -> b) -> ListT m a -> ListT m b)
-> (forall a b. a -> ListT m b -> ListT m a) -> Functor (ListT m)
forall a b. a -> ListT m b -> ListT m a
forall a b. (a -> b) -> ListT m a -> ListT m b
forall (m :: * -> *) a b. Functor m => a -> ListT m b -> ListT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> ListT m a -> ListT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> ListT m a -> ListT m b
fmap :: forall a b. (a -> b) -> ListT m a -> ListT m b
$c<$ :: forall (m :: * -> *) a b. Functor m => a -> ListT m b -> ListT m a
<$ :: forall a b. a -> ListT m b -> ListT m a
Functor, Functor (ListT m)
Functor (ListT m) =>
(forall a. a -> ListT m a)
-> (forall a b. ListT m (a -> b) -> ListT m a -> ListT m b)
-> (forall a b c.
    (a -> b -> c) -> ListT m a -> ListT m b -> ListT m c)
-> (forall a b. ListT m a -> ListT m b -> ListT m b)
-> (forall a b. ListT m a -> ListT m b -> ListT m a)
-> Applicative (ListT m)
forall a. a -> ListT m a
forall a b. ListT m a -> ListT m b -> ListT m a
forall a b. ListT m a -> ListT m b -> ListT m b
forall a b. ListT m (a -> b) -> ListT m a -> ListT m b
forall a b c. (a -> b -> c) -> ListT m a -> ListT m b -> ListT m 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 :: * -> *). Applicative m => Functor (ListT m)
forall (m :: * -> *) a. Applicative m => a -> ListT m a
forall (m :: * -> *) a b.
Applicative m =>
ListT m a -> ListT m b -> ListT m a
forall (m :: * -> *) a b.
Applicative m =>
ListT m a -> ListT m b -> ListT m b
forall (m :: * -> *) a b.
Applicative m =>
ListT m (a -> b) -> ListT m a -> ListT m b
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c) -> ListT m a -> ListT m b -> ListT m c
$cpure :: forall (m :: * -> *) a. Applicative m => a -> ListT m a
pure :: forall a. a -> ListT m a
$c<*> :: forall (m :: * -> *) a b.
Applicative m =>
ListT m (a -> b) -> ListT m a -> ListT m b
<*> :: forall a b. ListT m (a -> b) -> ListT m a -> ListT m b
$cliftA2 :: forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c) -> ListT m a -> ListT m b -> ListT m c
liftA2 :: forall a b c. (a -> b -> c) -> ListT m a -> ListT m b -> ListT m c
$c*> :: forall (m :: * -> *) a b.
Applicative m =>
ListT m a -> ListT m b -> ListT m b
*> :: forall a b. ListT m a -> ListT m b -> ListT m b
$c<* :: forall (m :: * -> *) a b.
Applicative m =>
ListT m a -> ListT m b -> ListT m a
<* :: forall a b. ListT m a -> ListT m b -> ListT m a
Applicative, Applicative (ListT m)
Applicative (ListT m) =>
(forall a. ListT m a)
-> (forall a. ListT m a -> ListT m a -> ListT m a)
-> (forall a. ListT m a -> ListT m [a])
-> (forall a. ListT m a -> ListT m [a])
-> Alternative (ListT m)
forall a. ListT m a
forall a. ListT m a -> ListT m [a]
forall a. ListT m a -> ListT m a -> ListT m a
forall (f :: * -> *).
Applicative f =>
(forall a. f a)
-> (forall a. f a -> f a -> f a)
-> (forall a. f a -> f [a])
-> (forall a. f a -> f [a])
-> Alternative f
forall (m :: * -> *). Alternative m => Applicative (ListT m)
forall (m :: * -> *) a. Alternative m => ListT m a
forall (m :: * -> *) a. Alternative m => ListT m a -> ListT m [a]
forall (m :: * -> *) a.
Alternative m =>
ListT m a -> ListT m a -> ListT m a
$cempty :: forall (m :: * -> *) a. Alternative m => ListT m a
empty :: forall a. ListT m a
$c<|> :: forall (m :: * -> *) a.
Alternative m =>
ListT m a -> ListT m a -> ListT m a
<|> :: forall a. ListT m a -> ListT m a -> ListT m a
$csome :: forall (m :: * -> *) a. Alternative m => ListT m a -> ListT m [a]
some :: forall a. ListT m a -> ListT m [a]
$cmany :: forall (m :: * -> *) a. Alternative m => ListT m a -> ListT m [a]
many :: forall a. ListT m a -> ListT m [a]
Alternative)

listT :: m [a] -> ListT m a
listT :: forall (m :: * -> *) a. m [a] -> ListT m a
listT = Compose m [] a -> ListT m a
forall (m :: * -> *) a. Compose m [] a -> ListT m a
ListT (Compose m [] a -> ListT m a)
-> (m [a] -> Compose m [] a) -> m [a] -> ListT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m [a] -> Compose m [] a
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose

runListT :: ListT m a -> m [a]
runListT :: forall (m :: * -> *) a. ListT m a -> m [a]
runListT = Compose m [] a -> m [a]
forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (Compose m [] a -> m [a])
-> (ListT m a -> Compose m [] a) -> ListT m a -> m [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ListT m a -> Compose m [] a
forall (m :: * -> *) a. ListT m a -> Compose m [] a
getListT

-- | This monad instance is _unlawful_,
-- it is only by accident and careful construction that it can be used here.
instance (Monad m) => Monad (ListT m) where
  ListT m a
ma >>= :: forall a b. ListT m a -> (a -> ListT m b) -> ListT m b
>>= a -> ListT m b
f = Compose m [] b -> ListT m b
forall (m :: * -> *) a. Compose m [] a -> ListT m a
ListT (Compose m [] b -> ListT m b) -> Compose m [] b -> ListT m b
forall a b. (a -> b) -> a -> b
$ m [b] -> Compose m [] b
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (m [b] -> Compose m [] b) -> m [b] -> Compose m [] b
forall a b. (a -> b) -> a -> b
$ do
    [a]
as <- ListT m a -> m [a]
forall (m :: * -> *) a. ListT m a -> m [a]
runListT ListT m a
ma
    ([[b]] -> [b]) -> m [[b]] -> m [b]
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[b]] -> [b]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (m [[b]] -> m [b]) -> m [[b]] -> m [b]
forall a b. (a -> b) -> a -> b
$ [a] -> (a -> m [b]) -> m [[b]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [a]
as ((a -> m [b]) -> m [[b]]) -> (a -> m [b]) -> m [[b]]
forall a b. (a -> b) -> a -> b
$ ListT m b -> m [b]
forall (m :: * -> *) a. ListT m a -> m [a]
runListT (ListT m b -> m [b]) -> (a -> ListT m b) -> a -> m [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ListT m b
f

instance MonadTrans ListT where
  lift :: forall (m :: * -> *) a. Monad m => m a -> ListT m a
lift = Compose m [] a -> ListT m a
forall (m :: * -> *) a. Compose m [] a -> ListT m a
ListT (Compose m [] a -> ListT m a)
-> (m a -> Compose m [] a) -> m a -> ListT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m [a] -> Compose m [] a
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (m [a] -> Compose m [] a)
-> (m a -> m [a]) -> m a -> Compose m [] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> [a]) -> m a -> m [a]
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> [a]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance (MonadIO m) => MonadIO (ListT m) where
  liftIO :: forall a. IO a -> ListT m a
liftIO = m a -> ListT m a
forall (m :: * -> *) a. Monad m => m a -> ListT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> ListT m a) -> (IO a -> m a) -> IO a -> ListT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO

instance (MonadDistribution m) => MonadDistribution (ListT m) where
  random :: ListT m Double
random = m Double -> ListT m Double
forall (m :: * -> *) a. Monad m => m a -> ListT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m Double
forall (m :: * -> *). MonadDistribution m => m Double
random
  bernoulli :: Double -> ListT m Bool
bernoulli = m Bool -> ListT m Bool
forall (m :: * -> *) a. Monad m => m a -> ListT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Bool -> ListT m Bool)
-> (Double -> m Bool) -> Double -> ListT m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> m Bool
forall (m :: * -> *). MonadDistribution m => Double -> m Bool
bernoulli
  categorical :: forall (v :: * -> *). Vector v Double => v Double -> ListT m Int
categorical = m Int -> ListT m Int
forall (m :: * -> *) a. Monad m => m a -> ListT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m Int -> ListT m Int)
-> (v Double -> m Int) -> v Double -> ListT m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v Double -> m Int
forall (v :: * -> *). Vector v Double => v Double -> m Int
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> m Int
categorical

instance (MonadFactor m) => MonadFactor (ListT m) where
  score :: Log Double -> ListT m ()
score = m () -> ListT m ()
forall (m :: * -> *) a. Monad m => m a -> ListT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ListT m ())
-> (Log Double -> m ()) -> Log Double -> ListT m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Log Double -> m ()
forall (m :: * -> *). MonadFactor m => Log Double -> m ()
score

instance (MonadMeasure m) => MonadMeasure (ListT m)

-- | A collection of weighted samples, or particles.
newtype PopulationT m a = PopulationT {forall (m :: * -> *) a. PopulationT m a -> WeightedT (ListT m) a
getPopulationT :: WeightedT (ListT m) a}
  deriving newtype ((forall a b. (a -> b) -> PopulationT m a -> PopulationT m b)
-> (forall a b. a -> PopulationT m b -> PopulationT m a)
-> Functor (PopulationT m)
forall a b. a -> PopulationT m b -> PopulationT m a
forall a b. (a -> b) -> PopulationT m a -> PopulationT m b
forall (m :: * -> *) a b.
Functor m =>
a -> PopulationT m b -> PopulationT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> PopulationT m a -> PopulationT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> PopulationT m a -> PopulationT m b
fmap :: forall a b. (a -> b) -> PopulationT m a -> PopulationT m b
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> PopulationT m b -> PopulationT m a
<$ :: forall a b. a -> PopulationT m b -> PopulationT m a
Functor, Functor (PopulationT m)
Functor (PopulationT m) =>
(forall a. a -> PopulationT m a)
-> (forall a b.
    PopulationT m (a -> b) -> PopulationT m a -> PopulationT m b)
-> (forall a b c.
    (a -> b -> c)
    -> PopulationT m a -> PopulationT m b -> PopulationT m c)
-> (forall a b.
    PopulationT m a -> PopulationT m b -> PopulationT m b)
-> (forall a b.
    PopulationT m a -> PopulationT m b -> PopulationT m a)
-> Applicative (PopulationT m)
forall a. a -> PopulationT m a
forall a b. PopulationT m a -> PopulationT m b -> PopulationT m a
forall a b. PopulationT m a -> PopulationT m b -> PopulationT m b
forall a b.
PopulationT m (a -> b) -> PopulationT m a -> PopulationT m b
forall a b c.
(a -> b -> c)
-> PopulationT m a -> PopulationT m b -> PopulationT m c
forall (m :: * -> *). Monad m => Functor (PopulationT m)
forall (m :: * -> *) a. Monad m => a -> PopulationT m a
forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> PopulationT m b -> PopulationT m a
forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> PopulationT m b -> PopulationT m b
forall (m :: * -> *) a b.
Monad m =>
PopulationT m (a -> b) -> PopulationT m a -> PopulationT m b
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c)
-> PopulationT m a -> PopulationT m b -> PopulationT m 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
$cpure :: forall (m :: * -> *) a. Monad m => a -> PopulationT m a
pure :: forall a. a -> PopulationT m a
$c<*> :: forall (m :: * -> *) a b.
Monad m =>
PopulationT m (a -> b) -> PopulationT m a -> PopulationT m b
<*> :: forall a b.
PopulationT m (a -> b) -> PopulationT m a -> PopulationT m b
$cliftA2 :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c)
-> PopulationT m a -> PopulationT m b -> PopulationT m c
liftA2 :: forall a b c.
(a -> b -> c)
-> PopulationT m a -> PopulationT m b -> PopulationT m c
$c*> :: forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> PopulationT m b -> PopulationT m b
*> :: forall a b. PopulationT m a -> PopulationT m b -> PopulationT m b
$c<* :: forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> PopulationT m b -> PopulationT m a
<* :: forall a b. PopulationT m a -> PopulationT m b -> PopulationT m a
Applicative, Applicative (PopulationT m)
Applicative (PopulationT m) =>
(forall a b.
 PopulationT m a -> (a -> PopulationT m b) -> PopulationT m b)
-> (forall a b.
    PopulationT m a -> PopulationT m b -> PopulationT m b)
-> (forall a. a -> PopulationT m a)
-> Monad (PopulationT m)
forall a. a -> PopulationT m a
forall a b. PopulationT m a -> PopulationT m b -> PopulationT m b
forall a b.
PopulationT m a -> (a -> PopulationT m b) -> PopulationT m b
forall (m :: * -> *). Monad m => Applicative (PopulationT m)
forall (m :: * -> *) a. Monad m => a -> PopulationT m a
forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> PopulationT m b -> PopulationT m b
forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> (a -> PopulationT m b) -> PopulationT m 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
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> (a -> PopulationT m b) -> PopulationT m b
>>= :: forall a b.
PopulationT m a -> (a -> PopulationT m b) -> PopulationT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
PopulationT m a -> PopulationT m b -> PopulationT m b
>> :: forall a b. PopulationT m a -> PopulationT m b -> PopulationT m b
$creturn :: forall (m :: * -> *) a. Monad m => a -> PopulationT m a
return :: forall a. a -> PopulationT m a
Monad, Monad (PopulationT m)
Monad (PopulationT m) =>
(forall a. IO a -> PopulationT m a) -> MonadIO (PopulationT m)
forall a. IO a -> PopulationT m a
forall (m :: * -> *).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
forall (m :: * -> *). MonadIO m => Monad (PopulationT m)
forall (m :: * -> *) a. MonadIO m => IO a -> PopulationT m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> PopulationT m a
liftIO :: forall a. IO a -> PopulationT m a
MonadIO, Monad (PopulationT m)
PopulationT m Double
Monad (PopulationT m) =>
PopulationT m Double
-> (Double -> Double -> PopulationT m Double)
-> (Double -> Double -> PopulationT m Double)
-> (Double -> Double -> PopulationT m Double)
-> (Double -> Double -> PopulationT m Double)
-> (Double -> PopulationT m Bool)
-> (forall (v :: * -> *).
    Vector v Double =>
    v Double -> PopulationT m Int)
-> (forall (v :: * -> *).
    (Vector v (Log Double), Vector v Double) =>
    v (Log Double) -> PopulationT m Int)
-> (forall a. [a] -> PopulationT m a)
-> (Double -> PopulationT m Int)
-> (Double -> PopulationT m Int)
-> (forall (v :: * -> *).
    Vector v Double =>
    v Double -> PopulationT m (v Double))
-> MonadDistribution (PopulationT m)
Double -> PopulationT m Bool
Double -> PopulationT m Int
Double -> Double -> PopulationT m Double
forall a. [a] -> PopulationT m a
forall (m :: * -> *).
Monad m =>
m Double
-> (Double -> Double -> m Double)
-> (Double -> Double -> m Double)
-> (Double -> Double -> m Double)
-> (Double -> Double -> m Double)
-> (Double -> m Bool)
-> (forall (v :: * -> *). Vector v Double => v Double -> m Int)
-> (forall (v :: * -> *).
    (Vector v (Log Double), Vector v Double) =>
    v (Log Double) -> m Int)
-> (forall a. [a] -> m a)
-> (Double -> m Int)
-> (Double -> m Int)
-> (forall (v :: * -> *).
    Vector v Double =>
    v Double -> m (v Double))
-> MonadDistribution m
forall (v :: * -> *).
Vector v Double =>
v Double -> PopulationT m (v Double)
forall (v :: * -> *).
Vector v Double =>
v Double -> PopulationT m Int
forall (v :: * -> *).
(Vector v (Log Double), Vector v Double) =>
v (Log Double) -> PopulationT m Int
forall (m :: * -> *). MonadDistribution m => Monad (PopulationT m)
forall (m :: * -> *). MonadDistribution m => PopulationT m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> PopulationT m Bool
forall (m :: * -> *).
MonadDistribution m =>
Double -> PopulationT m Int
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> PopulationT m Double
forall (m :: * -> *) a.
MonadDistribution m =>
[a] -> PopulationT m a
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> PopulationT m (v Double)
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> PopulationT m Int
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> PopulationT m Int
$crandom :: forall (m :: * -> *). MonadDistribution m => PopulationT m Double
random :: PopulationT m Double
$cuniform :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> PopulationT m Double
uniform :: Double -> Double -> PopulationT m Double
$cnormal :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> PopulationT m Double
normal :: Double -> Double -> PopulationT m Double
$cgamma :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> PopulationT m Double
gamma :: Double -> Double -> PopulationT m Double
$cbeta :: forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> PopulationT m Double
beta :: Double -> Double -> PopulationT m Double
$cbernoulli :: forall (m :: * -> *).
MonadDistribution m =>
Double -> PopulationT m Bool
bernoulli :: Double -> PopulationT m Bool
$ccategorical :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> PopulationT m Int
categorical :: forall (v :: * -> *).
Vector v Double =>
v Double -> PopulationT m Int
$clogCategorical :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> PopulationT m Int
logCategorical :: forall (v :: * -> *).
(Vector v (Log Double), Vector v Double) =>
v (Log Double) -> PopulationT m Int
$cuniformD :: forall (m :: * -> *) a.
MonadDistribution m =>
[a] -> PopulationT m a
uniformD :: forall a. [a] -> PopulationT m a
$cgeometric :: forall (m :: * -> *).
MonadDistribution m =>
Double -> PopulationT m Int
geometric :: Double -> PopulationT m Int
$cpoisson :: forall (m :: * -> *).
MonadDistribution m =>
Double -> PopulationT m Int
poisson :: Double -> PopulationT m Int
$cdirichlet :: forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> PopulationT m (v Double)
dirichlet :: forall (v :: * -> *).
Vector v Double =>
v Double -> PopulationT m (v Double)
MonadDistribution, Monad (PopulationT m)
Monad (PopulationT m) =>
(Log Double -> PopulationT m ()) -> MonadFactor (PopulationT m)
Log Double -> PopulationT m ()
forall (m :: * -> *). Monad m => Monad (PopulationT m)
forall (m :: * -> *). Monad m => Log Double -> PopulationT m ()
forall (m :: * -> *).
Monad m =>
(Log Double -> m ()) -> MonadFactor m
$cscore :: forall (m :: * -> *). Monad m => Log Double -> PopulationT m ()
score :: Log Double -> PopulationT m ()
MonadFactor, MonadFactor (PopulationT m)
MonadDistribution (PopulationT m)
(MonadDistribution (PopulationT m), MonadFactor (PopulationT m)) =>
MonadMeasure (PopulationT m)
forall (m :: * -> *).
MonadDistribution m =>
MonadFactor (PopulationT m)
forall (m :: * -> *).
MonadDistribution m =>
MonadDistribution (PopulationT m)
forall (m :: * -> *).
(MonadDistribution m, MonadFactor m) =>
MonadMeasure m
MonadMeasure)

instance MonadTrans PopulationT where
  lift :: forall (m :: * -> *) a. Monad m => m a -> PopulationT m a
lift = WeightedT (ListT m) a -> PopulationT m a
forall (m :: * -> *) a. WeightedT (ListT m) a -> PopulationT m a
PopulationT (WeightedT (ListT m) a -> PopulationT m a)
-> (m a -> WeightedT (ListT m) a) -> m a -> PopulationT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ListT m a -> WeightedT (ListT m) a
forall (m :: * -> *) a. Monad m => m a -> WeightedT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ListT m a -> WeightedT (ListT m) a)
-> (m a -> ListT m a) -> m a -> WeightedT (ListT m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> ListT m a
forall (m :: * -> *) a. Monad m => m a -> ListT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

-- | Explicit representation of the weighted sample with weights in the log
-- domain.
runPopulationT :: PopulationT m a -> m [(a, Log Double)]
runPopulationT :: forall (m :: * -> *) a. PopulationT m a -> m [(a, Log Double)]
runPopulationT = ListT m (a, Log Double) -> m [(a, Log Double)]
forall (m :: * -> *) a. ListT m a -> m [a]
runListT (ListT m (a, Log Double) -> m [(a, Log Double)])
-> (PopulationT m a -> ListT m (a, Log Double))
-> PopulationT m a
-> m [(a, Log Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WeightedT (ListT m) a -> ListT m (a, Log Double)
forall (m :: * -> *) a. WeightedT m a -> m (a, Log Double)
runWeightedT (WeightedT (ListT m) a -> ListT m (a, Log Double))
-> (PopulationT m a -> WeightedT (ListT m) a)
-> PopulationT m a
-> ListT m (a, Log Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PopulationT m a -> WeightedT (ListT m) a
forall (m :: * -> *) a. PopulationT m a -> WeightedT (ListT m) a
getPopulationT

-- | Explicit representation of the weighted sample.
explicitPopulation :: (Functor m) => PopulationT m a -> m [(a, Double)]
explicitPopulation :: forall (m :: * -> *) a.
Functor m =>
PopulationT m a -> m [(a, Double)]
explicitPopulation = ([(a, Log Double)] -> [(a, Double)])
-> m [(a, Log Double)] -> m [(a, Double)]
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((a, Log Double) -> (a, Double))
-> [(a, Log Double)] -> [(a, Double)]
forall a b. (a -> b) -> [a] -> [b]
map ((Log Double -> Double) -> (a, Log Double) -> (a, Double)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double)
-> (Log Double -> Double) -> Log Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Log Double -> Double
forall a. Log a -> a
ln))) (m [(a, Log Double)] -> m [(a, Double)])
-> (PopulationT m a -> m [(a, Log Double)])
-> PopulationT m a
-> m [(a, Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PopulationT m a -> m [(a, Log Double)]
forall (m :: * -> *) a. PopulationT m a -> m [(a, Log Double)]
runPopulationT

-- | Initialize 'PopulationT' with a concrete weighted sample.
fromWeightedList :: (Monad m) => m [(a, Log Double)] -> PopulationT m a
fromWeightedList :: forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> PopulationT m a
fromWeightedList = WeightedT (ListT m) a -> PopulationT m a
forall (m :: * -> *) a. WeightedT (ListT m) a -> PopulationT m a
PopulationT (WeightedT (ListT m) a -> PopulationT m a)
-> (m [(a, Log Double)] -> WeightedT (ListT m) a)
-> m [(a, Log Double)]
-> PopulationT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ListT m (a, Log Double) -> WeightedT (ListT m) a
forall (m :: * -> *) a.
Monad m =>
m (a, Log Double) -> WeightedT m a
weightedT (ListT m (a, Log Double) -> WeightedT (ListT m) a)
-> (m [(a, Log Double)] -> ListT m (a, Log Double))
-> m [(a, Log Double)]
-> WeightedT (ListT m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m [(a, Log Double)] -> ListT m (a, Log Double)
forall (m :: * -> *) a. m [a] -> ListT m a
listT

-- | Increase the sample size by a given factor.
-- The weights are adjusted such that their sum is preserved.
-- It is therefore safe to use 'spawn' in arbitrary places in the program
-- without introducing bias.
spawn :: (Monad m) => Int -> PopulationT m ()
spawn :: forall (m :: * -> *). Monad m => Int -> PopulationT m ()
spawn Int
n = m [((), Log Double)] -> PopulationT m ()
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> PopulationT m a
fromWeightedList (m [((), Log Double)] -> PopulationT m ())
-> m [((), Log Double)] -> PopulationT m ()
forall a b. (a -> b) -> a -> b
$ [((), Log Double)] -> m [((), Log Double)]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([((), Log Double)] -> m [((), Log Double)])
-> [((), Log Double)] -> m [((), Log Double)]
forall a b. (a -> b) -> a -> b
$ Int -> ((), Log Double) -> [((), Log Double)]
forall a. Int -> a -> [a]
replicate Int
n ((), Log Double
1 Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Int -> Log Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)

withParticles :: (Monad m) => Int -> PopulationT m a -> PopulationT m a
withParticles :: forall (m :: * -> *) a.
Monad m =>
Int -> PopulationT m a -> PopulationT m a
withParticles Int
n = (Int -> PopulationT m ()
forall (m :: * -> *). Monad m => Int -> PopulationT m ()
spawn Int
n PopulationT m () -> PopulationT m a -> PopulationT m a
forall a b. PopulationT m a -> PopulationT m b -> PopulationT m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>)

resampleGeneric ::
  (MonadDistribution m) =>
  -- | resampler
  (V.Vector Double -> m [Int]) ->
  PopulationT m a ->
  PopulationT m a
resampleGeneric :: forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> PopulationT m a -> PopulationT m a
resampleGeneric Vector Double -> m [Int]
resampler PopulationT m a
m = m [(a, Log Double)] -> PopulationT m a
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> PopulationT m a
fromWeightedList (m [(a, Log Double)] -> PopulationT m a)
-> m [(a, Log Double)] -> PopulationT m a
forall a b. (a -> b) -> a -> b
$ do
  [(a, Log Double)]
pop <- PopulationT m a -> m [(a, Log Double)]
forall (m :: * -> *) a. PopulationT m a -> m [(a, Log Double)]
runPopulationT PopulationT m a
m
  let ([a]
xs, [Log Double]
ps) = [(a, Log Double)] -> ([a], [Log Double])
forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  let n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
  let z :: Log Double
z = [Log Double] -> Log Double
forall a (f :: * -> *).
(RealFloat a, Foldable f) =>
f (Log a) -> Log a
Log.sum [Log Double]
ps
  if Log Double
z Log Double -> Log Double -> Bool
forall a. Ord a => a -> a -> Bool
> Log Double
0
    then do
      let weights :: Vector Double
weights = [Double] -> Vector Double
forall a. [a] -> Vector a
V.fromList ((Log Double -> Double) -> [Log Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double)
-> (Log Double -> Double) -> Log Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Log Double -> Double
forall a. Log a -> a
ln (Log Double -> Double)
-> (Log Double -> Log Double) -> Log Double -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Log Double
z)) [Log Double]
ps)
      [Int]
ancestors <- Vector Double -> m [Int]
resampler Vector Double
weights
      let xvec :: Vector a
xvec = [a] -> Vector a
forall a. [a] -> Vector a
V.fromList [a]
xs
      let offsprings :: [a]
offsprings = (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Vector a
xvec Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.!) [Int]
ancestors
      [(a, Log Double)] -> m [(a, Log Double)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, Log Double)] -> m [(a, Log Double)])
-> [(a, Log Double)] -> m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ (a -> (a, Log Double)) -> [a] -> [(a, Log Double)]
forall a b. (a -> b) -> [a] -> [b]
map (,Log Double
z Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Int -> Log Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) [a]
offsprings
    else -- if all weights are zero do not resample
      [(a, Log Double)] -> m [(a, Log Double)]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return [(a, Log Double)]
pop

-- | Systematic sampler.
-- Sample \(n\) values from \((0,1]\) as follows
-- \[
-- \begin{aligned}
-- u^{(1)} &\sim U\left(0, \frac{1}{n}\right] \\
-- u^{(i)} &=u^{(1)}+\frac{i-1}{n}, \quad i=2,3, \ldots, n
-- \end{aligned}
-- \]
-- and then pick integers \(m\) according to
-- \[
-- Q^{(m-1)}<u^{(n)} \leq Q^{(m)}
-- \]
-- where
-- \[
-- Q^{(m)}=\sum_{k=1}^{m} w^{(k)}
-- \]
-- and \(w^{(k)}\) are the weights. See also [Comparison of Resampling Schemes for Particle Filtering](https://arxiv.org/abs/cs/0507025).
systematic :: Double -> V.Vector Double -> [Int]
systematic :: Double -> Vector Double -> [Int]
systematic Double
u Vector Double
ps = Int -> Double -> Int -> Double -> [Int] -> [Int]
f Int
0 (Double
u Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) Int
0 Double
0 []
  where
    prob :: Int -> Double
prob Int
i = Vector Double
ps Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
V.! Int
i
    n :: Int
n = Vector Double -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Double
ps
    inc :: Double
inc = Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    f :: Int -> Double -> Int -> Double -> [Int] -> [Int]
f Int
i Double
_ Int
_ Double
_ [Int]
acc | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = [Int]
acc
    f Int
i Double
v Int
j Double
q [Int]
acc =
      if Double
v Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
q
        then Int -> Double -> Int -> Double -> [Int] -> [Int]
f (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
inc) Int
j Double
q (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
acc)
        else Int -> Double -> Int -> Double -> [Int] -> [Int]
f Int
i Double
v (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Double
q Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
prob Int
j) [Int]
acc

-- | Resample the population using the underlying monad and a systematic resampling scheme.
-- The total weight is preserved.
resampleSystematic ::
  (MonadDistribution m) =>
  PopulationT m a ->
  PopulationT m a
resampleSystematic :: forall (m :: * -> *) a.
MonadDistribution m =>
PopulationT m a -> PopulationT m a
resampleSystematic = (Vector Double -> m [Int]) -> PopulationT m a -> PopulationT m a
forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> PopulationT m a -> PopulationT m a
resampleGeneric (\Vector Double
ps -> (Double -> Vector Double -> [Int]
`systematic` Vector Double
ps) (Double -> [Int]) -> m Double -> m [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Double
forall (m :: * -> *). MonadDistribution m => m Double
random)

-- | Stratified sampler.
--
-- Sample \(n\) values from \((0,1]\) as follows
-- \[
-- u^{(i)} \sim U\left(\frac{i-1}{n}, \frac{i}{n}\right], \quad i=1,2, \ldots, n
-- \]
-- and then pick integers \(m\) according to
-- \[
-- Q^{(m-1)}<u^{(n)} \leq Q^{(m)}
-- \]
-- where
-- \[
-- Q^{(m)}=\sum_{k=1}^{m} w^{(k)}
-- \]
-- and \(w^{(k)}\) are the weights.
--
-- The conditional variance of stratified sampling is always smaller than that of multinomial sampling and it is also unbiased - see  [Comparison of Resampling Schemes for Particle Filtering](https://arxiv.org/abs/cs/0507025).
stratified :: (MonadDistribution m) => V.Vector Double -> m [Int]
stratified :: forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
stratified Vector Double
weights = do
  let bigN :: Int
bigN = Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
weights
  Vector Double
dithers <- Int -> m Double -> m (Vector Double)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM Int
bigN (Double -> Double -> m Double
forall (m :: * -> *).
MonadDistribution m =>
Double -> Double -> m Double
uniform Double
0.0 Double
1.0)
  let positions :: Vector Double
positions =
        (Double -> Double) -> Vector Double -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
bigN) (Vector Double -> Vector Double) -> Vector Double -> Vector Double
forall a b. (a -> b) -> a -> b
$
          (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Vector Double
dithers ((Int -> Double) -> Vector Int -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Int -> Vector Double) -> Vector Int -> Vector Double
forall a b. (a -> b) -> a -> b
$ [Int] -> Vector Int
forall a. [a] -> Vector a
V.fromList [Int
0 .. Int
bigN Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
      cumulativeSum :: Vector Double
cumulativeSum = (Double -> Double -> Double)
-> Double -> Vector Double -> Vector Double
forall a b. (a -> b -> a) -> a -> Vector b -> Vector a
V.scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0.0 Vector Double
weights
      coalg :: (Int, Int) -> Maybe (Maybe Int, (Int, Int))
coalg (Int
i, Int
j)
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
bigN =
            if (Vector Double
positions Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
! Int
i) Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< (Vector Double
cumulativeSum Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
! Int
j)
              then (Maybe Int, (Int, Int)) -> Maybe (Maybe Int, (Int, Int))
forall a. a -> Maybe a
Just (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
j, (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Int
j))
              else (Maybe Int, (Int, Int)) -> Maybe (Maybe Int, (Int, Int))
forall a. a -> Maybe a
Just (Maybe Int
forall a. Maybe a
Nothing, (Int
i, Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        | Bool
otherwise =
            Maybe (Maybe Int, (Int, Int))
forall a. Maybe a
Nothing
  [Int] -> m [Int]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int] -> m [Int]) -> [Int] -> m [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ [Maybe Int] -> [Int]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Int] -> [Int]) -> [Maybe Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ ((Int, Int) -> Maybe (Maybe Int, (Int, Int)))
-> (Int, Int) -> [Maybe Int]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr (Int, Int) -> Maybe (Maybe Int, (Int, Int))
coalg (Int
0, Int
0)

-- | Resample the population using the underlying monad and a stratified resampling scheme.
-- The total weight is preserved.
resampleStratified ::
  (MonadDistribution m) =>
  PopulationT m a ->
  PopulationT m a
resampleStratified :: forall (m :: * -> *) a.
MonadDistribution m =>
PopulationT m a -> PopulationT m a
resampleStratified = (Vector Double -> m [Int]) -> PopulationT m a -> PopulationT m a
forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> PopulationT m a -> PopulationT m a
resampleGeneric Vector Double -> m [Int]
forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
stratified

-- | Multinomial sampler.  Sample from \(0, \ldots, n - 1\) \(n\)
-- times drawn at random according to the weights where \(n\) is the
-- length of vector of weights.
multinomial :: (MonadDistribution m) => V.Vector Double -> m [Int]
multinomial :: forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
multinomial Vector Double
ps = Int -> m Int -> m [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
ps) (Vector Double -> m Int
forall (v :: * -> *). Vector v Double => v Double -> m Int
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v Double) =>
v Double -> m Int
categorical Vector Double
ps)

-- | Resample the population using the underlying monad and a multinomial resampling scheme.
-- The total weight is preserved.
resampleMultinomial ::
  (MonadDistribution m) =>
  PopulationT m a ->
  PopulationT m a
resampleMultinomial :: forall (m :: * -> *) a.
MonadDistribution m =>
PopulationT m a -> PopulationT m a
resampleMultinomial = (Vector Double -> m [Int]) -> PopulationT m a -> PopulationT m a
forall (m :: * -> *) a.
MonadDistribution m =>
(Vector Double -> m [Int]) -> PopulationT m a -> PopulationT m a
resampleGeneric Vector Double -> m [Int]
forall (m :: * -> *).
MonadDistribution m =>
Vector Double -> m [Int]
multinomial

-- | Separate the sum of weights into the 'WeightedT' transformer.
-- Weights are normalized after this operation.
extractEvidence ::
  (Monad m) =>
  PopulationT m a ->
  PopulationT (WeightedT m) a
extractEvidence :: forall (m :: * -> *) a.
Monad m =>
PopulationT m a -> PopulationT (WeightedT m) a
extractEvidence PopulationT m a
m = WeightedT m [(a, Log Double)] -> PopulationT (WeightedT m) a
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> PopulationT m a
fromWeightedList (WeightedT m [(a, Log Double)] -> PopulationT (WeightedT m) a)
-> WeightedT m [(a, Log Double)] -> PopulationT (WeightedT m) a
forall a b. (a -> b) -> a -> b
$ do
  [(a, Log Double)]
pop <- m [(a, Log Double)] -> WeightedT m [(a, Log Double)]
forall (m :: * -> *) a. Monad m => m a -> WeightedT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m [(a, Log Double)] -> WeightedT m [(a, Log Double)])
-> m [(a, Log Double)] -> WeightedT m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ PopulationT m a -> m [(a, Log Double)]
forall (m :: * -> *) a. PopulationT m a -> m [(a, Log Double)]
runPopulationT PopulationT m a
m
  let ([a]
xs, [Log Double]
ps) = [(a, Log Double)] -> ([a], [Log Double])
forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  let z :: Log Double
z = [Log Double] -> Log Double
forall a (f :: * -> *).
(RealFloat a, Foldable f) =>
f (Log a) -> Log a
sum [Log Double]
ps
  let ws :: [Log Double]
ws = (Log Double -> Log Double) -> [Log Double] -> [Log Double]
forall a b. (a -> b) -> [a] -> [b]
map (if Log Double
z Log Double -> Log Double -> Bool
forall a. Ord a => a -> a -> Bool
> Log Double
0 then (Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Log Double
z) else Log Double -> Log Double -> Log Double
forall a b. a -> b -> a
const (Log Double
1 Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Int -> Log Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Log Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Log Double]
ps))) [Log Double]
ps
  Log Double -> WeightedT m ()
forall (m :: * -> *). MonadFactor m => Log Double -> m ()
factor Log Double
z
  [(a, Log Double)] -> WeightedT m [(a, Log Double)]
forall a. a -> WeightedT m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, Log Double)] -> WeightedT m [(a, Log Double)])
-> [(a, Log Double)] -> WeightedT m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ [a] -> [Log Double] -> [(a, Log Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
xs [Log Double]
ws

-- | Push the evidence estimator as a score to the transformed monad.
-- Weights are normalized after this operation.
pushEvidence ::
  (MonadFactor m) =>
  PopulationT m a ->
  PopulationT m a
pushEvidence :: forall (m :: * -> *) a.
MonadFactor m =>
PopulationT m a -> PopulationT m a
pushEvidence = (forall x. WeightedT m x -> m x)
-> PopulationT (WeightedT m) a -> PopulationT m a
forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> PopulationT m a -> PopulationT n a
hoist WeightedT m x -> m x
forall x. WeightedT m x -> m x
forall (m :: * -> *) a. MonadFactor m => WeightedT m a -> m a
applyWeight (PopulationT (WeightedT m) a -> PopulationT m a)
-> (PopulationT m a -> PopulationT (WeightedT m) a)
-> PopulationT m a
-> PopulationT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PopulationT m a -> PopulationT (WeightedT m) a
forall (m :: * -> *) a.
Monad m =>
PopulationT m a -> PopulationT (WeightedT m) a
extractEvidence

-- | A properly weighted single sample, that is one picked at random according
-- to the weights, with the sum of all weights.
proper ::
  (MonadDistribution m) =>
  PopulationT m a ->
  WeightedT m a
proper :: forall (m :: * -> *) a.
MonadDistribution m =>
PopulationT m a -> WeightedT m a
proper PopulationT m a
m = do
  [(a, Log Double)]
pop <- PopulationT (WeightedT m) a -> WeightedT m [(a, Log Double)]
forall (m :: * -> *) a. PopulationT m a -> m [(a, Log Double)]
runPopulationT (PopulationT (WeightedT m) a -> WeightedT m [(a, Log Double)])
-> PopulationT (WeightedT m) a -> WeightedT m [(a, Log Double)]
forall a b. (a -> b) -> a -> b
$ PopulationT m a -> PopulationT (WeightedT m) a
forall (m :: * -> *) a.
Monad m =>
PopulationT m a -> PopulationT (WeightedT m) a
extractEvidence PopulationT m a
m
  let ([a]
xs, [Log Double]
ps) = [(a, Log Double)] -> ([a], [Log Double])
forall a b. [(a, b)] -> ([a], [b])
unzip [(a, Log Double)]
pop
  Int
index <- Vector (Log Double) -> WeightedT m Int
forall (v :: * -> *).
(Vector v (Log Double), Vector v Double) =>
v (Log Double) -> WeightedT m Int
forall (m :: * -> *) (v :: * -> *).
(MonadDistribution m, Vector v (Log Double), Vector v Double) =>
v (Log Double) -> m Int
logCategorical (Vector (Log Double) -> WeightedT m Int)
-> Vector (Log Double) -> WeightedT m Int
forall a b. (a -> b) -> a -> b
$ [Log Double] -> Vector (Log Double)
forall a. [a] -> Vector a
V.fromList [Log Double]
ps
  let x :: a
x = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
index
  a -> WeightedT m a
forall a. a -> WeightedT m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x

-- | Model evidence estimator, also known as pseudo-marginal likelihood.
evidence :: (Monad m) => PopulationT m a -> m (Log Double)
evidence :: forall (m :: * -> *) a.
Monad m =>
PopulationT m a -> m (Log Double)
evidence = WeightedT m [(a, Log Double)] -> m (Log Double)
forall (m :: * -> *) a.
Functor m =>
WeightedT m a -> m (Log Double)
extractWeight (WeightedT m [(a, Log Double)] -> m (Log Double))
-> (PopulationT m a -> WeightedT m [(a, Log Double)])
-> PopulationT m a
-> m (Log Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PopulationT (WeightedT m) a -> WeightedT m [(a, Log Double)]
forall (m :: * -> *) a. PopulationT m a -> m [(a, Log Double)]
runPopulationT (PopulationT (WeightedT m) a -> WeightedT m [(a, Log Double)])
-> (PopulationT m a -> PopulationT (WeightedT m) a)
-> PopulationT m a
-> WeightedT m [(a, Log Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PopulationT m a -> PopulationT (WeightedT m) a
forall (m :: * -> *) a.
Monad m =>
PopulationT m a -> PopulationT (WeightedT m) a
extractEvidence

-- | Picks one point from the population and uses model evidence as a 'score'
-- in the transformed monad.
-- This way a single sample can be selected from a population without
-- introducing bias.
collapse ::
  (MonadMeasure m) =>
  PopulationT m a ->
  m a
collapse :: forall (m :: * -> *) a. MonadMeasure m => PopulationT m a -> m a
collapse = WeightedT m a -> m a
forall (m :: * -> *) a. MonadFactor m => WeightedT m a -> m a
applyWeight (WeightedT m a -> m a)
-> (PopulationT m a -> WeightedT m a) -> PopulationT m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PopulationT m a -> WeightedT m a
forall (m :: * -> *) a.
MonadDistribution m =>
PopulationT m a -> WeightedT m a
proper

-- | PopulationT average of a function, computed using unnormalized weights.
popAvg :: (Monad m) => (a -> Double) -> PopulationT m a -> m Double
popAvg :: forall (m :: * -> *) a.
Monad m =>
(a -> Double) -> PopulationT m a -> m Double
popAvg a -> Double
f PopulationT m a
p = do
  [(a, Double)]
xs <- PopulationT m a -> m [(a, Double)]
forall (m :: * -> *) a.
Functor m =>
PopulationT m a -> m [(a, Double)]
explicitPopulation PopulationT m a
p
  let ys :: [Double]
ys = ((a, Double) -> Double) -> [(a, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\(a
x, Double
w) -> a -> Double
f a
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
w) [(a, Double)]
xs
  let t :: Double
t = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Data.List.sum [Double]
ys
  Double -> m Double
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Double
t

-- | Applies a transformation to the inner monad.
hoist ::
  (Monad n) =>
  (forall x. m x -> n x) ->
  PopulationT m a ->
  PopulationT n a
hoist :: forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> PopulationT m a -> PopulationT n a
hoist forall x. m x -> n x
f = n [(a, Log Double)] -> PopulationT n a
forall (m :: * -> *) a.
Monad m =>
m [(a, Log Double)] -> PopulationT m a
fromWeightedList (n [(a, Log Double)] -> PopulationT n a)
-> (PopulationT m a -> n [(a, Log Double)])
-> PopulationT m a
-> PopulationT n a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m [(a, Log Double)] -> n [(a, Log Double)]
forall x. m x -> n x
f (m [(a, Log Double)] -> n [(a, Log Double)])
-> (PopulationT m a -> m [(a, Log Double)])
-> PopulationT m a
-> n [(a, Log Double)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PopulationT m a -> m [(a, Log Double)]
forall (m :: * -> *) a. PopulationT m a -> m [(a, Log Double)]
runPopulationT