{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}

{- | Metropolis-Hastings inference.
-}

module Inference.MH (
    mh
  , mhStep
  , runMH
  , traceLPs
  , handleSamp
  , lookupSample
  , accept) where

import Control.Monad ( (>=>) )
import Data.Kind (Type)
import Data.Map (Map)
import Data.Maybe ( fromJust )
import Data.Set (Set, (\\))
import Effects.Dist ( Addr, Tag, Dist, Observe(..), Sample(..), pattern Obs, pattern Samp)
import Effects.ObsReader ( ObsReader )
import Effects.State ( State, modify, handleState )
import Env ( Env )
import Inference.SIM (handleObs, traceSamples)
import Model ( Model, handleCore )
import OpenSum (OpenSum(..))
import PrimDist
    ( ErasedPrimDist(ErasedPrimDist),
      PrimVal,
      PrimDist(UniformDist, DiscrUniformDist),
      pattern PrimDistPrf,
      sample )
import Prog ( Member(prj), EffectSum, Prog(..), discharge )
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified OpenSum
import Sampler ( Sampler, liftS )
import Trace ( LPTrace, FromSTrace(..), STrace, updateLPTrace )
import Unsafe.Coerce ( unsafeCoerce )

-- | Top-level wrapper for Metropolis-Hastings (MH) inference
mh :: (FromSTrace env, es ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe, Sample])
  -- | number of MH iterations
  => Int
  -- | model awaiting an input
  -> (b -> Model env es a)
  -- | (model input, input model environment)
  -> (b, Env env)
  -- | optional list of observable variable names (strings) to specify sample sites of interest
  {- For example, provide "mu" to specify interest in sampling #mu. This causes other variables to not be resampled unless necessary. -}
  -> [Tag]
  -- | [output model environment]
  -> Sampler [Env env]
mh :: forall (env :: [Assign Symbol (*)]) (es :: [* -> *]) b a.
(FromSTrace env,
 es
 ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
     Sample]) =>
Int
-> (b -> Model env es a)
-> (b, Env env)
-> [[Char]]
-> Sampler [Env env]
mh Int
n b -> Model env es a
model  (b
x_0, Env env
env_0) [[Char]]
tags = do
  -- Perform initial run of MH with no proposal sample site
  ((a, STrace), LPTrace)
y0 <- Env env
-> STrace
-> Addr
-> Model env es a
-> Sampler ((a, STrace), LPTrace)
forall (es :: [* -> *]) (env :: [Assign Symbol (*)]) a.
(es
 ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
     Sample]) =>
Env env
-> STrace
-> Addr
-> Model env es a
-> Sampler ((a, STrace), LPTrace)
runMH Env env
env_0 STrace
forall k a. Map k a
Map.empty ([Char]
"", Int
0) (b -> Model env es a
model b
x_0)
  -- Perform n MH iterations
  [((a, STrace), LPTrace)]
mhTrace <- (([((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)])
 -> ([((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)])
 -> [((a, STrace), LPTrace)]
 -> Sampler [((a, STrace), LPTrace)])
-> ([((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)])
-> [[((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)]]
-> [((a, STrace), LPTrace)]
-> Sampler [((a, STrace), LPTrace)]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ([((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)])
-> ([((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)])
-> [((a, STrace), LPTrace)]
-> Sampler [((a, STrace), LPTrace)]
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
(>=>) [((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)]
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
-> ([((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)])
-> [[((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)]]
forall a. Int -> a -> [a]
replicate Int
n (Env env
-> Model env es a
-> [[Char]]
-> [((a, STrace), LPTrace)]
-> Sampler [((a, STrace), LPTrace)]
forall (es :: [* -> *]) (env :: [Assign Symbol (*)]) a.
(es
 ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
     Sample]) =>
Env env
-> Model env es a
-> [[Char]]
-> [((a, STrace), LPTrace)]
-> Sampler [((a, STrace), LPTrace)]
mhStep Env env
env_0 (b -> Model env es a
model b
x_0) [[Char]]
tags)) [((a, STrace), LPTrace)
y0]
  -- Return sample trace
  [Env env] -> Sampler [Env env]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Env env] -> Sampler [Env env]) -> [Env env] -> Sampler [Env env]
forall a b. (a -> b) -> a -> b
$ (((a, STrace), LPTrace) -> Env env)
-> [((a, STrace), LPTrace)] -> [Env env]
forall a b. (a -> b) -> [a] -> [b]
map (\((a
_, STrace
strace), LPTrace
_) -> STrace -> Env env
forall (env :: [Assign Symbol (*)]).
FromSTrace env =>
STrace -> Env env
fromSTrace STrace
strace) [((a, STrace), LPTrace)]
mhTrace

-- | Perform one step of MH
mhStep :: (es ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe, Sample])
  -- | model environment
  => Env env
  -- | model
  -> Model env es a
  -- | tags indicating sample sites of interest
  -> [Tag]
  -- | trace of previous MH outputs
  -> [((a, STrace), LPTrace)]
  -- | updated trace of MH outputs
  -> Sampler [((a, STrace), LPTrace)]
mhStep :: forall (es :: [* -> *]) (env :: [Assign Symbol (*)]) a.
(es
 ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
     Sample]) =>
Env env
-> Model env es a
-> [[Char]]
-> [((a, STrace), LPTrace)]
-> Sampler [((a, STrace), LPTrace)]
mhStep Env env
env Model env es a
model [[Char]]
tags [((a, STrace), LPTrace)]
trace = do
  -- Get previous mh output
  let ((a
x, STrace
samples), LPTrace
logps) = [((a, STrace), LPTrace)] -> ((a, STrace), LPTrace)
forall a. [a] -> a
head [((a, STrace), LPTrace)]
trace
  -- Get possible addresses to propose new samples for
      sampleSites :: STrace
sampleSites = if [[Char]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Char]]
tags then STrace
samples
                    else  (Addr -> (ErasedPrimDist, OpenSum PrimVal) -> Bool)
-> STrace -> STrace
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (\([Char]
tag, Int
i) (ErasedPrimDist, OpenSum PrimVal)
_ -> [Char]
tag [Char] -> [[Char]] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [[Char]]
tags) STrace
samples
  -- Draw a proposal sample address
  Int
α_samp_ind <- PrimDist Int -> Sampler Int
forall a. PrimDist a -> Sampler a
sample (PrimDist Int -> Sampler Int) -> PrimDist Int -> Sampler Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> PrimDist Int
DiscrUniformDist Int
0 (STrace -> Int
forall k a. Map k a -> Int
Map.size STrace
sampleSites Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  let (Addr
α_samp, (ErasedPrimDist, OpenSum PrimVal)
_) = Int -> STrace -> (Addr, (ErasedPrimDist, OpenSum PrimVal))
forall k a. Int -> Map k a -> (k, a)
Map.elemAt Int
α_samp_ind STrace
sampleSites
  -- Run MH with proposal sample address
  ((a
x', STrace
samples'), LPTrace
logps') <- Env env
-> STrace
-> Addr
-> Model env es a
-> Sampler ((a, STrace), LPTrace)
forall (es :: [* -> *]) (env :: [Assign Symbol (*)]) a.
(es
 ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
     Sample]) =>
Env env
-> STrace
-> Addr
-> Model env es a
-> Sampler ((a, STrace), LPTrace)
runMH Env env
env STrace
samples Addr
α_samp Model env es a
model
  -- Compute acceptance ratio
  Double
acceptance_ratio <- IO Double -> Sampler Double
forall a. IO a -> Sampler a
liftS (IO Double -> Sampler Double) -> IO Double -> Sampler Double
forall a b. (a -> b) -> a -> b
$ Addr -> STrace -> STrace -> LPTrace -> LPTrace -> IO Double
accept Addr
α_samp STrace
samples STrace
samples' LPTrace
logps LPTrace
logps'
  Double
u <- PrimDist Double -> Sampler Double
forall a. PrimDist a -> Sampler a
sample (Double -> Double -> PrimDist Double
UniformDist Double
0 Double
1)
  if Double
u Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
acceptance_ratio
    then do [((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)]
forall (m :: * -> *) a. Monad m => a -> m a
return (((a
x', STrace
samples'), LPTrace
logps')((a, STrace), LPTrace)
-> [((a, STrace), LPTrace)] -> [((a, STrace), LPTrace)]
forall a. a -> [a] -> [a]
:[((a, STrace), LPTrace)]
trace)
    else do [((a, STrace), LPTrace)] -> Sampler [((a, STrace), LPTrace)]
forall (m :: * -> *) a. Monad m => a -> m a
return [((a, STrace), LPTrace)]
trace

-- | Handler for one iteration of MH
runMH :: (es ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe, Sample])
  -- | model environment
  => Env env
  -- | sample trace of previous MH iteration
  -> STrace
  -- | sample address of interest
  -> Addr
  -- | model
  -> Model env es a
  -- | (model output, sample trace, log-probability trace)
  -> Sampler ((a, STrace), LPTrace)
runMH :: forall (es :: [* -> *]) (env :: [Assign Symbol (*)]) a.
(es
 ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
     Sample]) =>
Env env
-> STrace
-> Addr
-> Model env es a
-> Sampler ((a, STrace), LPTrace)
runMH Env env
env STrace
strace Addr
α_samp =
     STrace
-> Addr
-> Prog '[Sample] ((a, STrace), LPTrace)
-> Sampler ((a, STrace), LPTrace)
forall a. STrace -> Addr -> Prog '[Sample] a -> Sampler a
handleSamp STrace
strace Addr
α_samp  (Prog '[Sample] ((a, STrace), LPTrace)
 -> Sampler ((a, STrace), LPTrace))
-> (Model
      env
      '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
        Sample]
      a
    -> Prog '[Sample] ((a, STrace), LPTrace))
-> Model
     env
     '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
       Sample]
     a
-> Sampler ((a, STrace), LPTrace)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog '[Observe, Sample] ((a, STrace), LPTrace)
-> Prog '[Sample] ((a, STrace), LPTrace)
forall (es :: [* -> *]) a. Prog (Observe : es) a -> Prog es a
handleObs
   (Prog '[Observe, Sample] ((a, STrace), LPTrace)
 -> Prog '[Sample] ((a, STrace), LPTrace))
-> (Model
      env
      '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
        Sample]
      a
    -> Prog '[Observe, Sample] ((a, STrace), LPTrace))
-> Model
     env
     '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
       Sample]
     a
-> Prog '[Sample] ((a, STrace), LPTrace)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LPTrace
-> Prog '[State LPTrace, Observe, Sample] (a, STrace)
-> Prog '[Observe, Sample] ((a, STrace), LPTrace)
forall s (es :: [* -> *]) a.
s -> Prog (State s : es) a -> Prog es (a, s)
handleState LPTrace
forall k a. Map k a
Map.empty (Prog '[State LPTrace, Observe, Sample] (a, STrace)
 -> Prog '[Observe, Sample] ((a, STrace), LPTrace))
-> (Model
      env
      '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
        Sample]
      a
    -> Prog '[State LPTrace, Observe, Sample] (a, STrace))
-> Model
     env
     '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
       Sample]
     a
-> Prog '[Observe, Sample] ((a, STrace), LPTrace)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STrace
-> Prog '[State STrace, State LPTrace, Observe, Sample] a
-> Prog '[State LPTrace, Observe, Sample] (a, STrace)
forall s (es :: [* -> *]) a.
s -> Prog (State s : es) a -> Prog es (a, s)
handleState STrace
forall k a. Map k a
Map.empty
   (Prog '[State STrace, State LPTrace, Observe, Sample] a
 -> Prog '[State LPTrace, Observe, Sample] (a, STrace))
-> (Model
      env
      '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
        Sample]
      a
    -> Prog '[State STrace, State LPTrace, Observe, Sample] a)
-> Model
     env
     '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
       Sample]
     a
-> Prog '[State LPTrace, Observe, Sample] (a, STrace)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog '[State STrace, State LPTrace, Observe, Sample] a
-> Prog '[State STrace, State LPTrace, Observe, Sample] a
forall (es :: [* -> *]) a.
(Member (State LPTrace) es, Member Sample es, Member Observe es) =>
Prog es a -> Prog es a
traceLPs (Prog '[State STrace, State LPTrace, Observe, Sample] a
 -> Prog '[State STrace, State LPTrace, Observe, Sample] a)
-> (Model
      env
      '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
        Sample]
      a
    -> Prog '[State STrace, State LPTrace, Observe, Sample] a)
-> Model
     env
     '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
       Sample]
     a
-> Prog '[State STrace, State LPTrace, Observe, Sample] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prog '[State STrace, State LPTrace, Observe, Sample] a
-> Prog '[State STrace, State LPTrace, Observe, Sample] a
forall (es :: [* -> *]) a.
(Member (State STrace) es, Member Sample es) =>
Prog es a -> Prog es a
traceSamples (Prog '[State STrace, State LPTrace, Observe, Sample] a
 -> Prog '[State STrace, State LPTrace, Observe, Sample] a)
-> (Model
      env
      '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
        Sample]
      a
    -> Prog '[State STrace, State LPTrace, Observe, Sample] a)
-> Model
     env
     '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
       Sample]
     a
-> Prog '[State STrace, State LPTrace, Observe, Sample] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Env env
-> Model
     env
     '[ObsReader env, Dist, State STrace, State LPTrace, Observe,
       Sample]
     a
-> Prog '[State STrace, State LPTrace, Observe, Sample] a
forall (es :: [* -> *]) (env :: [Assign Symbol (*)]) a.
(Member Observe es, Member Sample es) =>
Env env -> Model env (ObsReader env : Dist : es) a -> Prog es a
handleCore Env env
env




-- | Handler for tracing log-probabilities for each @Sample@ and @Observe@ operation
traceLPs :: (Member (State LPTrace) es, Member Sample es, Member Observe es)
  => Prog es a
  -> Prog es a
traceLPs :: forall (es :: [* -> *]) a.
(Member (State LPTrace) es, Member Sample es, Member Observe es) =>
Prog es a -> Prog es a
traceLPs (Val a
x) = a -> Prog es a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
traceLPs (Op EffectSum es x
op x -> Prog es a
k) = case EffectSum es x
op of
  Samp (PrimDistPrf PrimDist x
d) Addr
α ->
       EffectSum es x -> (x -> Prog es a) -> Prog es a
forall (es :: [* -> *]) x a.
EffectSum es x -> (x -> Prog es a) -> Prog es a
Op EffectSum es x
op (\x
x -> (LPTrace -> LPTrace) -> Prog es ()
forall s (es :: [* -> *]).
Member (State s) es =>
(s -> s) -> Prog es ()
modify (Addr -> PrimDist x -> x -> LPTrace -> LPTrace
forall x. Addr -> PrimDist x -> x -> LPTrace -> LPTrace
updateLPTrace Addr
α PrimDist x
d x
x) Prog es () -> Prog es a -> Prog es a
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
                    Prog es a -> Prog es a
forall (es :: [* -> *]) a.
(Member (State LPTrace) es, Member Sample es, Member Observe es) =>
Prog es a -> Prog es a
traceLPs (x -> Prog es a
k x
x))
  Obs PrimDist x
d x
y Addr
α ->
    EffectSum es x -> (x -> Prog es a) -> Prog es a
forall (es :: [* -> *]) x a.
EffectSum es x -> (x -> Prog es a) -> Prog es a
Op EffectSum es x
op (\ x
x -> (LPTrace -> LPTrace) -> Prog es ()
forall s (es :: [* -> *]).
Member (State s) es =>
(s -> s) -> Prog es ()
modify (Addr -> PrimDist x -> x -> LPTrace -> LPTrace
forall x. Addr -> PrimDist x -> x -> LPTrace -> LPTrace
updateLPTrace Addr
α PrimDist x
d x
y) Prog es () -> Prog es a -> Prog es a
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
                  Prog es a -> Prog es a
forall (es :: [* -> *]) a.
(Member (State LPTrace) es, Member Sample es, Member Observe es) =>
Prog es a -> Prog es a
traceLPs (x -> Prog es a
k x
y))
  EffectSum es x
_         -> EffectSum es x -> (x -> Prog es a) -> Prog es a
forall (es :: [* -> *]) x a.
EffectSum es x -> (x -> Prog es a) -> Prog es a
Op EffectSum es x
op (Prog es a -> Prog es a
forall (es :: [* -> *]) a.
(Member (State LPTrace) es, Member Sample es, Member Observe es) =>
Prog es a -> Prog es a
traceLPs (Prog es a -> Prog es a) -> (x -> Prog es a) -> x -> Prog es a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> Prog es a
k)

-- | Handler for @Sample@ that selectively reuses old samples or draws new ones
handleSamp ::
  -- | sample trace
     STrace
  -- | address of the proposal sample site for the current MH iteration
  -> Addr
  -> Prog '[Sample] a
  -> Sampler a
handleSamp :: forall a. STrace -> Addr -> Prog '[Sample] a -> Sampler a
handleSamp STrace
strace Addr
α_samp (Op EffectSum '[Sample] x
op x -> Prog '[Sample] a
k) = case EffectSum '[Sample] x -> Either (EffectSum '[] x) (Sample x)
forall (e :: * -> *) (es :: [* -> *]) x.
EffectSum (e : es) x -> Either (EffectSum es x) (e x)
discharge EffectSum '[Sample] x
op of
  Right (Sample (PrimDistPrf PrimDist x
d) Addr
α) ->
        do x
x <- STrace -> PrimDist x -> Addr -> Addr -> Sampler x
forall a.
Member a PrimVal =>
STrace -> PrimDist a -> Addr -> Addr -> Sampler a
lookupSample STrace
strace PrimDist x
d Addr
α Addr
α_samp
           STrace -> Addr -> Prog '[Sample] a -> Sampler a
forall a. STrace -> Addr -> Prog '[Sample] a -> Sampler a
handleSamp STrace
strace Addr
α_samp (x -> Prog '[Sample] a
k x
x)
  Either (EffectSum '[] x) (Sample x)
_  -> [Char] -> Sampler a
forall a. HasCallStack => [Char] -> a
error [Char]
"Impossible: Nothing cannot occur"
handleSamp STrace
_ Addr
_ (Val a
x) = a -> Sampler a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x

-- | For a given address, look up a sampled value from a sample trace, returning
--   it only if the primitive distribution it was sampled from matches the current one.
lookupSample :: OpenSum.Member a PrimVal
  =>
  -- | sample trace
     STrace
  -- | distribution to sample from
  -> PrimDist a
  -- | address of current sample site
  -> Addr
  -- | address of proposal sample site
  -> Addr
  -> Sampler a
lookupSample :: forall a.
Member a PrimVal =>
STrace -> PrimDist a -> Addr -> Addr -> Sampler a
lookupSample STrace
samples PrimDist a
d Addr
α Addr
α_samp
  | Addr
α Addr -> Addr -> Bool
forall a. Eq a => a -> a -> Bool
== Addr
α_samp = PrimDist a -> Sampler a
forall a. PrimDist a -> Sampler a
sample PrimDist a
d
  | Bool
otherwise   =
      case Addr -> STrace -> Maybe (ErasedPrimDist, OpenSum PrimVal)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Addr
α STrace
samples of
        Just (ErasedPrimDist PrimDist a
d', OpenSum PrimVal
x) -> do
          if PrimDist a
d PrimDist a -> PrimDist a -> Bool
forall a. Eq a => a -> a -> Bool
== PrimDist a -> PrimDist a
forall a b. a -> b
unsafeCoerce PrimDist a
d'
            then a -> Sampler a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> a
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe a -> a) -> Maybe a -> a
forall a b. (a -> b) -> a -> b
$ OpenSum PrimVal -> Maybe a
forall a (as :: [*]). Member a as => OpenSum as -> Maybe a
OpenSum.prj OpenSum PrimVal
x)
            else PrimDist a -> Sampler a
forall a. PrimDist a -> Sampler a
sample PrimDist a
d
        Maybe (ErasedPrimDist, OpenSum PrimVal)
Nothing -> PrimDist a -> Sampler a
forall a. PrimDist a -> Sampler a
sample PrimDist a
d

-- | Compute acceptance probability
accept ::
  -- | address of new sampled value
     Addr
  -- | previous MH sample trace
  -> STrace
  -- | new MH sample trace
  -> STrace
  -- | previous MH log-probability trace
  -> LPTrace
  -- | current MH log-probability trace
  -> LPTrace
  -> IO Double
accept :: Addr -> STrace -> STrace -> LPTrace -> LPTrace -> IO Double
accept Addr
x0 STrace
_Ⲭ STrace
_Ⲭ' LPTrace
logℙ LPTrace
logℙ' = do
  let _X'sampled :: Set Addr
_X'sampled = Addr -> Set Addr
forall a. a -> Set a
Set.singleton Addr
x0 Set Addr -> Set Addr -> Set Addr
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` (STrace -> Set Addr
forall k a. Map k a -> Set k
Map.keysSet STrace
_Ⲭ' Set Addr -> Set Addr -> Set Addr
forall a. Ord a => Set a -> Set a -> Set a
\\ STrace -> Set Addr
forall k a. Map k a -> Set k
Map.keysSet STrace
_Ⲭ)
      _Xsampled :: Set Addr
_Xsampled  = Addr -> Set Addr
forall a. a -> Set a
Set.singleton Addr
x0 Set Addr -> Set Addr -> Set Addr
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` (STrace -> Set Addr
forall k a. Map k a -> Set k
Map.keysSet STrace
_Ⲭ Set Addr -> Set Addr -> Set Addr
forall a. Ord a => Set a -> Set a -> Set a
\\ STrace -> Set Addr
forall k a. Map k a -> Set k
Map.keysSet STrace
_Ⲭ')
  let dom_logα :: Double
dom_logα   = Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ STrace -> Int
forall k a. Map k a -> Int
Map.size STrace
_Ⲭ) Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> Double
forall a. Floating a => a -> a
log (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ STrace -> Int
forall k a. Map k a -> Int
Map.size STrace
_Ⲭ')
  let _Xlogα :: Double
_Xlogα     = (Double -> Addr -> Double) -> Double -> Set Addr -> Double
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Double
logα Addr
v -> Double
logα Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Maybe Double -> Double
forall a. HasCallStack => Maybe a -> a
fromJust (Addr -> LPTrace -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Addr
v LPTrace
logℙ))
                         Double
0 (LPTrace -> Set Addr
forall k a. Map k a -> Set k
Map.keysSet LPTrace
logℙ Set Addr -> Set Addr -> Set Addr
forall a. Ord a => Set a -> Set a -> Set a
\\ Set Addr
_Xsampled)
  let _X'logα :: Double
_X'logα    = (Double -> Addr -> Double) -> Double -> Set Addr -> Double
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Double
logα Addr
v -> Double
logα Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Maybe Double -> Double
forall a. HasCallStack => Maybe a -> a
fromJust (Addr -> LPTrace -> Maybe Double
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Addr
v LPTrace
logℙ'))
                         Double
0 (LPTrace -> Set Addr
forall k a. Map k a -> Set k
Map.keysSet LPTrace
logℙ' Set Addr -> Set Addr -> Set Addr
forall a. Ord a => Set a -> Set a -> Set a
\\ Set Addr
_X'sampled)
  Double -> IO Double
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> IO Double) -> Double -> IO Double
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
exp (Double
dom_logα Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
_X'logα Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
_Xlogα)