-- | Sequential Monte Carlo on a discrete hidden Markov model.
--
-- A 2-state HMM with binary observations.  Exact inference by enumeration
-- of all 2^T state sequences.  SMC particle filter approximates the same
-- posterior.  Oracle: L1 distance between SMC and exact marginal < 0.1.
module Circuit.Inference.SMC
  ( -- * Model
    State (..),
    transProb,
    obsProb,
    initialP,

    -- * Exact inference
    exactFiltering,

    -- * SMC particle filter
    particleFilter,

    -- * Polynomial shape oracle
    SMCPoly,
    smcIn,
    smcSystem,
    smcTotalWeight,

    -- * Oracle
    trace5,
    l1Distance,
  )
where

import Circuit.Poly (Dir, Mono, Poly (..))
import Circuit.Prob (Prob (..))
import Circuit.System (System, monoIn, runSystem, system)
import Data.List (foldl')
import Data.Maybe (fromMaybe, listToMaybe)
import Data.Void (Void, absurd)
import System.Random (StdGen, mkStdGen, randomR)

-- ---------------------------------------------------------------------------
-- 2-state HMM
-- ---------------------------------------------------------------------------

data State = StateA | StateB
  deriving (State -> State -> Bool
(State -> State -> Bool) -> (State -> State -> Bool) -> Eq State
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: State -> State -> Bool
== :: State -> State -> Bool
$c/= :: State -> State -> Bool
/= :: State -> State -> Bool
Eq, Eq State
Eq State =>
(State -> State -> Ordering)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> Bool)
-> (State -> State -> State)
-> (State -> State -> State)
-> Ord State
State -> State -> Bool
State -> State -> Ordering
State -> State -> State
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
$ccompare :: State -> State -> Ordering
compare :: State -> State -> Ordering
$c< :: State -> State -> Bool
< :: State -> State -> Bool
$c<= :: State -> State -> Bool
<= :: State -> State -> Bool
$c> :: State -> State -> Bool
> :: State -> State -> Bool
$c>= :: State -> State -> Bool
>= :: State -> State -> Bool
$cmax :: State -> State -> State
max :: State -> State -> State
$cmin :: State -> State -> State
min :: State -> State -> State
Ord, Int -> State -> ShowS
[State] -> ShowS
State -> String
(Int -> State -> ShowS)
-> (State -> String) -> ([State] -> ShowS) -> Show State
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> State -> ShowS
showsPrec :: Int -> State -> ShowS
$cshow :: State -> String
show :: State -> String
$cshowList :: [State] -> ShowS
showList :: [State] -> ShowS
Show)

allStates :: [State]
allStates :: [State]
allStates = [State
StateA, State
StateB]

-- | Initial distribution: P(S₁=S).
initialP :: State -> Double
initialP :: State -> Double
initialP State
StateA = Double
0.6
initialP State
StateB = Double
0.4

-- | Transition: P(next | current).
transProb :: State -> State -> Double
transProb :: State -> State -> Double
transProb State
StateA State
StateA = Double
0.7
transProb State
StateA State
StateB = Double
0.3
transProb State
StateB State
StateA = Double
0.2
transProb State
StateB State
StateB = Double
0.8

-- | Observation likelihood: P(o | state). o ∈ {0, 1}.
obsProb :: Int -> State -> Double
obsProb :: Int -> State -> Double
obsProb Int
1 State
StateA = Double
0.9
obsProb Int
1 State
StateB = Double
0.2
obsProb Int
0 State
StateA = Double
0.1
obsProb Int
0 State
StateB = Double
0.8
obsProb Int
_ State
_ = String -> Double
forall a. HasCallStack => String -> a
error String
"obsProb: observation must be 0 or 1"

-- | 5-step trace, 2⁵ = 32 sequences for exact enumeration.
trace5 :: [Int]
trace5 :: [Int]
trace5 = [Int
1, Int
0, Int
1, Int
0, Int
1]

-- ---------------------------------------------------------------------------
-- Exact inference by full enumeration
-- ---------------------------------------------------------------------------

-- | Compute joint probability P(S_1…S_T, O_1…O_T) for a given state sequence.
jointProb :: [State] -> [Int] -> Double
jointProb :: [State] -> [Int] -> Double
jointProb [State
s] [Int
o] = State -> Double
initialP State
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> State -> Double
obsProb Int
o State
s
jointProb (State
s : ss :: [State]
ss@(State
s' : [State]
_)) (Int
o : [Int]
os) =
  State -> Double
initialP State
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> State -> Double
obsProb Int
o State
s Double -> Double -> Double
forall a. Num a => a -> a -> a
* State -> [(State, Int)] -> Double
go State
s' ([State] -> [Int] -> [(State, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [State]
ss [Int]
os)
  where
    go :: State -> [(State, Int)] -> Double
go State
_ [] = Double
1
    go State
prev ((State
cur, Int
o') : [(State, Int)]
rest) = State -> State -> Double
transProb State
prev State
cur Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> State -> Double
obsProb Int
o' State
cur Double -> Double -> Double
forall a. Num a => a -> a -> a
* State -> [(State, Int)] -> Double
go State
cur [(State, Int)]
rest
jointProb [State]
_ [Int]
_ = String -> Double
forall a. HasCallStack => String -> a
error String
"jointProb: length mismatch"

-- | All state sequences of length t.
allSeqs :: Int -> [[State]]
allSeqs :: Int -> [[State]]
allSeqs Int
0 = [[]]
allSeqs Int
k = [State
s State -> [State] -> [State]
forall a. a -> [a] -> [a]
: [State]
rest | State
s <- [State]
allStates, [State]
rest <- Int -> [[State]]
allSeqs (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]

-- | Exact filtering P(S_T | O_{1:T}) by enumeration.
-- Returns [(S, normalised probability)].
exactFiltering :: [Int] -> [(State, Double)]
exactFiltering :: [Int] -> [(State, Double)]
exactFiltering [Int]
obs =
  let t :: Int
t = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
obs
      seqs :: [[State]]
seqs = Int -> [[State]]
allSeqs Int
t
      joints :: [(State, Double)]
joints = ([State] -> (State, Double)) -> [[State]] -> [(State, Double)]
forall a b. (a -> b) -> [a] -> [b]
map (\[State]
stateSeq -> ([State] -> State
forall a. HasCallStack => [a] -> a
last [State]
stateSeq, [State] -> [Int] -> Double
jointProb [State]
stateSeq [Int]
obs)) [[State]]
seqs
      -- Sum probability for each final state
      totalByState :: State -> Double
totalByState State
s = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double
p | (State
lastS, Double
p) <- [(State, Double)]
joints, State
lastS State -> State -> Bool
forall a. Eq a => a -> a -> Bool
== State
s]
      total :: Double
total = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((State, Double) -> Double) -> [(State, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (State, Double) -> Double
forall a b. (a, b) -> b
snd [(State, Double)]
joints)
   in [(State
s, State -> Double
totalByState State
s Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
total) | Double
total Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0, State
s <- [State]
allStates]

-- ---------------------------------------------------------------------------
-- Particle filter (SMC) — deterministic seed for reproducibility
-- ---------------------------------------------------------------------------

-- | Run a particle filter and return empirical final filtering distribution.
-- n particles, deterministic RNG (seed = 42), systematic resampling.
particleFilter :: [Int] -> Int -> [(State, Double)]
particleFilter :: [Int] -> Int -> [(State, Double)]
particleFilter [Int]
obs Int
nP =
  let gen0 :: StdGen
gen0 = Int -> StdGen
mkStdGen Int
42

      step :: ([State], StdGen) -> Int -> ([State], StdGen)
step ([State]
particles, StdGen
gen) Int
o =
        let -- Propagate each particle by sampling next state
            nextState :: State -> b -> (State, b)
nextState State
s b
g =
              let (Double
r, b
g') = (Double, Double) -> b -> (Double, b)
forall g. RandomGen g => (Double, Double) -> g -> (Double, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Double
0 :: Double, Double
1) b
g
               in (if Double
r Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< State -> State -> Double
transProb State
s State
StateA then State
StateA else State
StateB, b
g')
            ([State]
propagated, StdGen
gen1) = (([State], StdGen) -> State -> ([State], StdGen))
-> ([State], StdGen) -> [State] -> ([State], StdGen)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\([State]
ps, StdGen
g) State
p -> let (State
p', StdGen
g') = State -> StdGen -> (State, StdGen)
forall {b}. RandomGen b => State -> b -> (State, b)
nextState State
p StdGen
g in ([State]
ps [State] -> [State] -> [State]
forall a. [a] -> [a] -> [a]
++ [State
p'], StdGen
g')) ([], StdGen
gen) [State]
particles
            -- Weight by observation likelihood
            ws :: [Double]
ws = (State -> Double) -> [State] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\State
p -> Int -> State -> Double
obsProb Int
o State
p) [State]
propagated
            total :: Double
total = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
ws
            normWs :: [Double]
normWs = if Double
total Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 then (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
total) [Double]
ws else Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
nP (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
nP)
            -- Systematic resample
            ([State]
resampled, StdGen
gen2) = [(State, Double)] -> Int -> StdGen -> ([State], StdGen)
resample ([State] -> [Double] -> [(State, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [State]
propagated [Double]
normWs) Int
nP StdGen
gen1
         in ([State]
resampled, StdGen
gen2)

      -- Initial: all particles start in StateA (deterministic)
      initParticles :: [State]
initParticles = Int -> State -> [State]
forall a. Int -> a -> [a]
replicate Int
nP State
StateA

      ([State]
finalParticles, StdGen
_) = (([State], StdGen) -> Int -> ([State], StdGen))
-> ([State], StdGen) -> [Int] -> ([State], StdGen)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ([State], StdGen) -> Int -> ([State], StdGen)
step ([State]
initParticles, StdGen
gen0) [Int]
obs

      -- Filtering estimate: fraction in each state
      prob :: State -> Double
prob State
s = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([State] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((State -> Bool) -> [State] -> [State]
forall a. (a -> Bool) -> [a] -> [a]
filter (State -> State -> Bool
forall a. Eq a => a -> a -> Bool
== State
s) [State]
finalParticles)) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nP
   in [(State
s, State -> Double
prob State
s) | State
s <- [State]
allStates]

-- | Systematic resampling with a supplied RNG.
resample :: [(State, Double)] -> Int -> StdGen -> ([State], StdGen)
resample :: [(State, Double)] -> Int -> StdGen -> ([State], StdGen)
resample [(State, Double)]
weighted Int
n StdGen
gen =
  let cumWeights :: [Double]
cumWeights = Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
drop Int
1 ((Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0 (((State, Double) -> Double) -> [(State, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (State, Double) -> Double
forall a b. (a, b) -> b
snd [(State, Double)]
weighted))
      stepVal :: Double
stepVal = Double
1.0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
      (Double
u0, StdGen
gen') = (Double, Double) -> StdGen -> (Double, StdGen)
forall g. RandomGen g => (Double, Double) -> g -> (Double, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Double
0, Double
stepVal) StdGen
gen
      threshold :: Int -> Double
threshold Int
i = Double
u0 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
stepVal
      pick :: Double -> State
pick Double
u =
        State -> Maybe State -> State
forall a. a -> Maybe a -> a
fromMaybe (String -> State
forall a. HasCallStack => String -> a
error String
"pick: no particle above threshold") (Maybe State -> State) -> Maybe State -> State
forall a b. (a -> b) -> a -> b
$
          [State] -> Maybe State
forall a. [a] -> Maybe a
listToMaybe ([State] -> Maybe State) -> [State] -> Maybe State
forall a b. (a -> b) -> a -> b
$
            ((State, Double) -> State) -> [(State, Double)] -> [State]
forall a b. (a -> b) -> [a] -> [b]
map (State, Double) -> State
forall a b. (a, b) -> a
fst (((State, Double) -> Bool) -> [(State, Double)] -> [(State, Double)]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (\(State
_, Double
c) -> Double
c Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
u) ([State] -> [Double] -> [(State, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip (((State, Double) -> State) -> [(State, Double)] -> [State]
forall a b. (a -> b) -> [a] -> [b]
map (State, Double) -> State
forall a b. (a, b) -> a
fst [(State, Double)]
weighted) [Double]
cumWeights))
      picked :: [State]
picked = (Int -> State) -> [Int] -> [State]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> State
pick (Double -> State) -> (Int -> Double) -> Int -> State
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Double
threshold) [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
   in ([State]
picked, StdGen
gen')

-- ---------------------------------------------------------------------------
-- Polynomial shape oracle: SMC as Prod (Mono () State) (Const Double)
-- ---------------------------------------------------------------------------

-- | SMC polynomial: particle stream paired with a weight annotation.
--
-- * @Mono () State@ — the particle is an output position; the input direction
--   is trivial (@()@) because the particle is sampled, not supplied.
-- * @Const Double@ — the weight is an output position with /no/ direction,
--   so it cannot be fed back as an input.
--
-- Written directly in 'Poly' constructors because 'Mono' is a type synonym.
type SMCPoly = 'Prod ('Prod ('Const State) ('Exp ())) ('Const Double)

-- | Canonical input direction for the SMC polynomial: advance one step.
smcIn :: () -> Dir SMCPoly
smcIn :: () -> Dir SMCPoly
smcIn () = Either Void () -> Either (Either Void ()) Void
forall a b. a -> Either a b
Left (() -> Either Void ()
forall a b. b -> Either a b
Right ())

-- | One-step SMC kernel for a fixed observation.
--
-- Given current hidden state @s@ and observation @o@, transition to @s'@
-- according to @transProb@ and emit the particle @s'@ together with weight
-- @obsProb o s'@.  The weight is part of the output position, not an input
-- direction, which is exactly the instance-table claim.
smcSystem :: Int -> System (Prob (->) Double) State SMCPoly
smcSystem :: Int -> System (Prob (->) Double) State SMCPoly
smcSystem Int
o = Prob (->) Double (State, Dir SMCPoly) (State, Pos SMCPoly)
-> System (Prob (->) Double) State SMCPoly
forall (arr :: * -> * -> *) s (p :: Poly).
arr (s, Dir p) (s, Pos p) -> System arr s p
system (Prob (->) Double (State, Dir SMCPoly) (State, Pos SMCPoly)
 -> System (Prob (->) Double) State SMCPoly)
-> Prob (->) Double (State, Dir SMCPoly) (State, Pos SMCPoly)
-> System (Prob (->) Double) State SMCPoly
forall a b. (a -> b) -> a -> b
$ (forall x.
 ((x, (State, Pos SMCPoly)) -> Double)
 -> (x, (State, Dir SMCPoly)) -> Double)
-> Prob (->) Double (State, Dir SMCPoly) (State, Pos SMCPoly)
forall {k} (arr :: * -> k -> *) (r :: k) a b.
(forall x. arr (x, b) r -> arr (x, a) r) -> Prob arr r a b
Prob ((forall x.
  ((x, (State, Pos SMCPoly)) -> Double)
  -> (x, (State, Dir SMCPoly)) -> Double)
 -> Prob (->) Double (State, Dir SMCPoly) (State, Pos SMCPoly))
-> (forall x.
    ((x, (State, Pos SMCPoly)) -> Double)
    -> (x, (State, Dir SMCPoly)) -> Double)
-> Prob (->) Double (State, Dir SMCPoly) (State, Pos SMCPoly)
forall a b. (a -> b) -> a -> b
$ \(x, (State, Pos SMCPoly)) -> Double
k (x
x, (State
s, Dir SMCPoly
d)) ->
  case Dir SMCPoly
d of
    Left Either Void ()
dMono -> case Either Void ()
dMono of
      Left Void
v -> Void -> Double
forall a. Void -> a
absurd Void
v
      Right () ->
        (Double -> Double -> Double) -> Double -> [Double] -> Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
          Double -> Double -> Double
forall a. Num a => a -> a -> a
(+)
          Double
0
          [ State -> State -> Double
transProb State
s State
s' Double -> Double -> Double
forall a. Num a => a -> a -> a
* (x, (State, Pos SMCPoly)) -> Double
k (x
x, (State
s', ((State
s', ()), Int -> State -> Double
obsProb Int
o State
s')))
          | State
s' <- [State]
allStates
          ]
    Right Void
v -> Void -> Double
forall a. Void -> a
absurd Void
v

-- | Total expected weight emitted by one SMC step starting from @s@ for
-- observation @o@.  This is the marginal likelihood @P(o | s)@.
smcTotalWeight :: Int -> State -> Double
smcTotalWeight :: Int -> State -> Double
smcTotalWeight Int
o State
s =
  Prob
  (->)
  Double
  (State, Either (Either Void ()) Void)
  (State, ((State, ()), Double))
-> forall x.
   ((x, (State, ((State, ()), Double))) -> Double)
   -> (x, (State, Either (Either Void ()) Void)) -> Double
forall {k} (arr :: * -> k -> *) (r :: k) a b.
Prob arr r a b -> forall x. arr (x, b) r -> arr (x, a) r
runProb
    (System (Prob (->) Double) State SMCPoly
-> Prob (->) Double (State, Dir SMCPoly) (State, Pos SMCPoly)
forall (arr :: * -> * -> *) s (p :: Poly).
System arr s p -> arr (s, Dir p) (s, Pos p)
runSystem (Int -> System (Prob (->) Double) State SMCPoly
smcSystem Int
o))
    (\(()
_, (State
_s', ((State
_p, ()), Double
w))) -> Double
w)
    ((), (State
s, () -> Dir SMCPoly
smcIn ()))

-- ---------------------------------------------------------------------------
-- Oracle
-- ---------------------------------------------------------------------------

-- | L1 distance between two filtering distributions.
l1Distance :: [(State, Double)] -> [(State, Double)] -> Double
l1Distance :: [(State, Double)] -> [(State, Double)] -> Double
l1Distance [(State, Double)]
d1 [(State, Double)]
d2 =
  [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double -> Double
forall a. Num a => a -> a
abs (State -> [(State, Double)] -> Double
forall {c} {a}. (Num c, Eq a) => a -> [(a, c)] -> c
look State
s [(State, Double)]
d1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- State -> [(State, Double)] -> Double
forall {c} {a}. (Num c, Eq a) => a -> [(a, c)] -> c
look State
s [(State, Double)]
d2) | State
s <- [State]
allStates]
  where
    look :: a -> [(a, c)] -> c
look a
s = c -> Maybe c -> c
forall a. a -> Maybe a -> a
fromMaybe c
0 (Maybe c -> c) -> ([(a, c)] -> Maybe c) -> [(a, c)] -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [(a, c)] -> Maybe c
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup a
s