Enumerate.hs

-- | Exact credal inference by enumerating Knightian valuations.
--
--   The functions returning bounds throw an error if the evidence is
--   unsatisfiable, i.e. the credal set is empty.
module Imp.Inference.Enumerate
  ( preciseMarginal
  , credalVertices
  , intervalProbability
  , intervalExpectation
  , marginal
  ) where

import Control.Monad (foldM)
import Control.Monad.State.Strict (runState)
import Data.Bool (bool)
import Data.Functor.Compose (Compose(..))
import Data.Maybe (mapMaybe)
import qualified Data.IntMap.Strict as IntMap
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import Imp.BDD (VarLabel(..))
import Imp.BDD.Builder (bddRestrict)
import Imp.BDD.Compile (compile)
import Imp.BDD.WMC (Weight(..), wmcBatch)
import Imp.DSL (Imp)
import Imp.Semiring (ProbS(..))

-- | Flips get Bernoulli branch pair. Knights have been conditioned away.
probWeights :: Weight -> (ProbS, ProbS)
probWeights w = case w of
  Prob p   -> (ProbS (1 - p), ProbS p)
  Knight _ -> error "probWeights: Knights should have been conditioned away"

-- | Compute the min and max of a non-empty list in a single strict pass.
bounds :: Ord a => [a] -> (a, a)
bounds []     = error "No feasible probabilities"
bounds (x:xs) = foldl' step (x, x) xs
  where step (!mn, !mx) y = (min mn y, max mx y)

-- | Precise marginal distribution. Only for programs with no Knightian choices.
--   Outcomes with probability zero are omitted, unlike 'marginal', which reports
--   every return value.
preciseMarginal :: Ord a => Imp '[] a -> Map a Double
preciseMarginal prog =
  let (mgr, weights, _, worlds) = compile prog
      counts = unProb <$> wmcBatch probWeights mgr weights worlds
      total  = sum counts
  in if total > 0
       then Map.filter (> 0) ((/ total) <$> counts)
       else error "No feasible probabilities"

-- | Compute the vertices of the credal set, one distribution per
--   feasible Knightian valuation.  The extreme points are a subset
--   of these.
credalVertices :: Ord a => Imp g a -> [Map a Double]
credalVertices prog =
  let (mgr, weights, _, worlds) = compile prog
      valuations = sequence [ [(VarLabel k, False), (VarLabel k, True)]
                            | (k, Knight _) <- IntMap.toList weights ]
      condition val =
        traverse (\bdd -> foldM (\g (vl, b) -> bddRestrict g vl b) bdd val) worlds
      (valEvents, mgr') = runState (mapM condition valuations) mgr
      counts = unProb <$> wmcBatch probWeights mgr' weights (Compose valEvents)
      normalize row =
        let !total = sum row
        in if total > 0 then Just ((/ total) <$> row) else Nothing
  in mapMaybe normalize (getCompose counts)

-- | Exact lower and upper probability of an event, by valuation enumeration.
intervalProbability :: Ord a => Imp g a -> (a -> Bool) -> (Double, Double)
intervalProbability prog predicate =
  intervalExpectation prog (bool 0 1 . predicate)

-- | Exact lower and upper expectation of a real-valued function, by
--   valuation enumeration.
intervalExpectation :: Ord a => Imp g a -> (a -> Double) -> (Double, Double)
intervalExpectation prog score =
  bounds [ sum [score v * p | (v, p) <- Map.toList d]
         | d <- credalVertices prog ]

-- | Exact per-value marginal bounds, by valuation enumeration.
marginal :: Ord a => Imp g a -> Map a (Double, Double)
marginal prog = case credalVertices prog of
  []    -> error "No feasible probabilities"
  dists -> bounds <$> Map.unionsWith (++) [ (\x -> [x]) <$> d | d <- dists ]